Title: Offline Oracle-Efficient Learning for Contextual MDPs via Layerwise Exploration-Exploitation Tradeoff
Venue: NeurIPS 2024
Accepted: yes
Keywords: reinforcement learning

contextual MDP

offline density estimation

computational efficiency
Abstract: Motivated by the recent discovery of a statistical and computational reduction from contextual bandits to offline regression \citep{simchi2020bypassing}, we address the general (stochastic) Contextual Markov Decision Process (CMDP) problem with horizon $H$ (as known as CMDP with $H$ layers). In this paper, we introduce a reduction from CMDPs to offline density estimation under the realizability assumption, i.e., a model class $\mathcal{M}$ containing the true underlying CMDP is provided in advance. We develop an efficient, statistically near-optimal algorithm requiring only $O(H \log T)$ calls to an offline density estimation algorithm (or oracle) across all $T$ rounds. This number can be further reduced to $O(H \log \log T)$ if $T$ is known in advance. Our results mark the first efficient and near-optimal reduction from CMDPs to offline density estimation without imposing any structural assumptions on the model class. A notable feature of our algorithm is the design of a layerwise exploration-exploitation tradeoff tailored to address the layerwise structure of CMDPs. Additionally, our algorithm is versatile and applicable to pure exploration tasks in reward-free reinforcement learning.

Offline Oracle-Efficient Learning for Contextual MDPs via Layerwise Exploration-Exploitation Tradeoff

Abstract

Motivated by the recent discovery of a statistical and computational reduction from contextual bandits to offline regression [36], we address the general (stochastic) Contextual Markov Decision Process (CMDP) problem with horizon H (as known as CMDP with H layers). In this paper, we introduce a reduction from CMDPs to offline density estimation under the realizability assumption, i.e., a model class M containing the true underlying CMDP is provided in advance. We develop an efficient, statistically near-optimal algorithm requiring only O ( H log T ) calls to an offline density estimation algorithm (or oracle) across all T rounds of interaction. This number can be further reduced to O ( H log log T ) if T is known in advance. Our results mark the first efficient and near-optimal reduction from CMDPs to offline density estimation without imposing any structural assumptions on the model class. A notable feature of our algorithm is the design of a layerwise exploration-exploitation tradeoff tailored to address the layerwise structure of CMDPs. Additionally, our algorithm is versatile and applicable to pure exploration tasks in reward-free reinforcement learning.

1 Introduction

Markov Decision Processes (MDPs) model the long-term interaction between a learner and the environment and are used in diverse areas such as inventory management, recommendation systems, advertising, and healthcare [35, 37]. The Contextual MDP (CMDP) extends MDPs by incorporating external factors, known as contexts , such as gender, age, location, or device in customer interactions, or lab data and medical history in healthcare [19, 33]. In an H -layer CMDP, the learner receives an instantaneous reward at each step over H steps and aims to maximize the cumulative reward (return). For T rounds of interaction, the learner's performance is measured by regret , which is the difference between the total return obtained and that of an optimal policy.

In the special case of contextual bandits (one-layer CMDPs), a decade of research has led to algorithms with optimal regret bounds and efficient implementations with access to an offline regression algorithm (also termed as an offline regression oracle ) [13, 2, 3, 15, 14, 36, 40]. Most notably, Simchi-Levi and Xu [36] demonstrates an offline-oracle-based algorithm FALCON that achieves optimal regret for general (stochastic) contextual bandits with access to an offline regression oracle (e.g., the Empirical Risk Minimization (ERM) oracle). Moreover, the algorithm is efficient given the output of the offline oracle (also referred to as offline oracle-efficient) and requires only O (log log T ) calls to the oracle across all T rounds if T is known. These properties are highly desirable in practice since they reduce the computational problem of contextual bandits to the classical problem of offline regression with little overhead. However, to the best of our knowledge, no algorithm with these properties is available in the literature for general (stochastic) CMDPs. So, in this paper, we study the following question: Is there an offline-oracle-efficient algorithm that achieves the optimal regret for general (stochastic) CMDPs with only O ( H log log T ) number of oracle calls?

MIT IDSS

Several works have provided partial results for this question. Foster et al. [16] provides a general reduction from interactive decision making to online density estimation and has CMDP as an application. The proposed E2D algorithm achieves optimal regret but is online-oracle-efficient (as opposed to offline-oracle-efficient) since it requires access to an online density estimation algorithm. Foster et al. [18] provides a further reduction from online density estimation to offline density estimation, with the caveat that the reduction itself is inefficient. A similar online-oracle-efficient algorithm is developed by Levy et al. [26]. A separate thread of optimism-based algorithms for CMDPs extending the UCCB algorithm for contextual bandits [40] is studied by Levy and Mansour [24], Deng et al. [11] with either assumption on the reachability of the CMDP or the representation structure of the CMDP (see Appendix B for more details). Last but not least, the algorithms proposed by Foster et al. [16, 18], Levy and Mansour [24], Deng et al. [11] all require O ( T ) number of oracle calls to the online/offline oracle respectively.

In this work, we present an affirmative answer to the question by introducing the algorithm of LOLIPOP (Algorithm 1). For S number of states, A number of actions, and a given model class M where the underlying true model lies, the algorithm achieves the regret guarantee of ˜ O (poly( H,S,A ) √ T log |M| ) . This regret guarantee is minimax optimal up to poly( H,S,A ) factor [24]. The LOLIPOP algorithm assumes access to a Maximum Likelihood Estimation (MLE) oracle and is offline-oracle efficient. The results can be generalized to general offline density estimation oracles. The most notable technical features are: (1) It generalizes the FALCON algorithm by SimchiLevi and Xu [36] to adapt to the multi-layer structure of a CMDP. More specifically, the FALCON algorithm is divided into O (log log T ) epochs, each corresponding to an oracle call, a fixed randomized policy. However, it is known for the MDPs that the learner has to switch its randomized policy at least ˜ Ω( H ) times to achieve sublinear regret [44]. Indeed, we further divide each epoch into H segments, each with an oracle call, a new randomized policy for layerwise exploration-exploitation tradeoff. (2) In each segment, the exploration-exploitation tradeoff is done through Inverse Gap Weighting (IGW) on estimated regret for a set of explorative policies. The idea of running IGW on such a policy cover is proposed by Foster et al. [16]. However, their policy cover is designed for H -layer exploration-exploration tradeoff and only works with strong online estimation oracles. In contrast, we refine the estimation of the occupancy measure layerwise by introducing the trusted occupancy measures . This refinement enables our algorithm to work with offline oracles. (3) Many other policy cover-based methods [12, 29, 30, 5] are developed for exploration tasks. Most notably, Mhammedi et al. [29] clips the occupancy measures on states with low reachability. Our approach takes a step forward to clip all transitions with low reachability to compute the trusted occupancy measures.

Besides all the above novelties, the LOLIPOP algorithm is versatile and applicable to the pure exploration task of reward-free reinforcement learning for CMDPs. Concretely, it obtains nearoptimal sample complexity of O ( H 7 S 4 A 3 log( |M| /δ ) /ε 2 ) with only O ( H ) number of oracle calls. Both the sample complexity bound and computational efficiency result for reward-free learning for stochastic CMDPs are new to the best of our knowledge.

2 Preliminaries

We defer the standard notation and related works to Appendices A and B.

2.1 Problem Setup

A Contextual Makovian Decision Process (CMDP) is defined by the tuple ( C , M = { M ( c ) } c ∈C , S , A , s 1 ) , where C is the contextual space, S is the state space, A is the action space and s 1 ∈ S is a fixed starting state independent of the context. We focus on tabular CMDPs which assumes S = |S| , A = |A| < ∞ . For any context c ∈ C , M ( c ) = { P h M ( c ) , R h M ( c ) } H h =1 consists of H -layers of probability transition kernel { P h M ( c ) } H h =1 and reward distributions { R h M ( c ) } H h =1 , where P h M ( c ) and R h M ( c ) are specified by P h M ( · | s, a ; c ) ∈ ∆( S ) and R h M ( s, a ; c ) ∈ ∆([0 , 1]) for all h ∈ [ H ] and s, a ∈ S × A . For simplicity, we also denote M = { P h M , R h M } H h =1 , where P h M = { P h M ( c ) } c ∈C and R h M = { R h M ( c ) } c ∈C . Let Π RNS denote the set of all randomized, non-stationary policies, where any π = ( π 1 , . . . , π H ) ∈ Π RNS has π h : S → ∆( A ) for any h ∈ [ H ] . We use T to denote the total number of rounds, H to denote the horizon (the total number of layers). Let M ⋆ = { P h ⋆ , R h ⋆ } h ∈ [ H ] be the underlying true CMDP the learner interact with. The interactive protocal proceeds in T rounds, where for each round t , the t -th trajectory is generated as:

A context c t is draw i.i.d. from an unknown distribution D and s 1 t = s 1 .

The learner chooses the policy π t based on the context.

For h = 1 , . . . , H :

The action is drawn from the randomized policy a h t ∼ π h t ( s h t ) .

The reward and the next state is drawn respectively from the reward distribution and the transition kernel, i.e., r h t ∼ R h ⋆ ( s h t , a h t ; c t ) and s h +1 t ∼ P h ⋆ ( · | s h t , a h t ; c t ) .

Without lose of generality, throughout the paper, we assume that the total reward 0 ≤ ∑ H h =1 r h ≤ 1 almost surely. For any model M , context c and policy π , we use M ( π, c ) to denote the distribution of the trajectory c 1 , π 1 , s 1 1 , a 1 1 , r 1 1 , . . . , s H 1 , a H 1 , r H 1 given M ⋆ = M , c 1 = c , and π 1 = π . Also denote the probability and the expectation under M ( π, c ) to be P M,π,c ( · ) and E M,π,c [ · ] respectively. Given any policy π , state s and action a , we define the action value function Q h ⋆ ( s, a ; π, c ) at layer h and the value function V h ⋆ ( s ; π, c ) at layer h under context c and policy π as

We denote the optimal policy under context c as π ⋆,c and abbreviate its value function as V h ⋆ ( · ; c ) . For h = 1 , we further simply the notation by denoting V 1 ⋆ ( c ) = V 1 ⋆ ( s 1 ; c ) and V 1 ⋆ ( π, c ) = V 1 ⋆ ( s 1 ; π, c ) . The regret of policy π under context c and the total regret 1 of the learner are defined as

where E t [ · ] is the conditional expectation given the interaction up to round t .

Assumption 2.1 (Realizability) . The learner is given a model class M where the true underlying model M ⋆ lies, that is, M ⋆ ∈ M .

2.2 Offline Densitiy Estimation Oracles

For any model class M , a general offline density estimation oracle associated with M , denoted by OffDE M , is defined as an algorithm that generates a predictor ̂ M based on the input data and M . In this paper, we measure the performance of the predictor in terms of the squared Hellinger distance, which is defined for any two distributions P and Q for any common dominating measure ν 2 by

Using squared Hellinger distance for reinforcement learning is popularized by Foster et al. [16], and we adopt such a divergence for our purpose as well. Concretely, we are interested in the following statistical guarantee.

Definition 2.1 (Offline density estimation oracle) . Let p be a map from a context to a distribution on the set of policies Π RNS , that is, for any c ∈ C , p ( c ) ∈ ∆(Π RNS ) . Given n training trajectories ( c i , π i , s 1 i , a 1 i , r 1 i , . . . , s H i , a H i , r H i ) i.i.d. drawn according to c i ∼ D , π i ∼ p ( c i ) and

s 1 i , a 1 i , r 1 i , . . . , s H i , a H i , r H i be the trajectory sampled according to M ⋆ ( π i , c i ) . The offline density estimation oracle OffDE M returns a predictor ̂ M . For any δ ∈ (0 , 1 / 2) , with probability at least 1 -δ , we have

The Maximum Likelihood Estimation estimator MLE M is an example of an offline density estimation oracle that achieves E M ,δ ( n ) ≲ log( |M| /δ ) /n (see more details in Appendix C) and can be implemented using ERM on the log loss. Moreover, this implementation can be efficient for cases like the multinomial logit model [32, Example 2.4] and [1].

3 Main Results and Algorithm

In this section, we present our main results and introduce the algorithm of LOLIPOP (Algorithm 1). First, we give an overview of the algorithm. Then, we discuss the theoretical guarantees obtained by this algorithm. Finally, we introduce the algorithm's different components with corresponding guarantees. All proofs are deferred to Appendix D.

3.1 Main Results

Overview of Algorithm 1. The algorithm proceeds with epochs. The total number of T rounds is divided into N epochs. For an epoch schedule 0 = τ 0 < τ 1 < · · · < τ N = T/H to be specified later, the m -th epoch will last H ( τ m -τ m -1 ) rounds. Furthermore, each epoch is evenly divided into H segments, each consisting of τ m -τ m -1 rounds. During the h -th segment in m -th epoch, a kernel p h m : C → ∆(Π) will be specified to determine the policy. More specifically, upon receiving the context c t , a policy π t will be sampled from p h m ( c t ) and executed. After collecting the trajectories { c t , π t } ∪ { s j t , a j t , r j t } j ∈ [ H ] in the h -th segment of the m -th epoch for τ m -1 H +( τ m -τ m -1 )( h -1) + 1 ≤ t ≤ τ m -1 H +( τ m -τ m -1 ) h , the offline density estimation oracle OffDE M is called with these trajectories as input. Denote the output ̂ M h m , we will only be interested in the h -th layer of this output, which we denote by { ̂ P h m , ̂ R h m } . Then the collections of estimators ̂ M m = { ̂ P h m , ̂ R h m } H h =1 will be used for the next epoch.

Throughout this paper, we will adopt the following convention for the free variables m,π,c,h,s, a . They will be used to denote an epoch index in [ N ] , a policy in Π RNS , a context in C , a layer index in [ H ] , a state in S , and an action in A respectively.

Before we dive into the details of the algorithm, we highlight first the theoretical guarantees obtained.

Theorem 3.1. If T is known, then by choosing the epoch schedule τ m = 2( T/H ) 1 -2 -m for m ≥ 1 and the offline density estimation oracle OffDE M = MLE M , the outputs { π t } t ∈ [ T ] of Algorithm 1 satisfies that with probability at least 1 -δ ,

with only O ( H log log T ) number of oracle calls to the MLE M oracle for δ ∈ (0 , 1 / 2) .

Theorem 3.2. If T is not known, then by choosing the epoch schedule τ m = 2 m for m ≥ 1 and the offline density estimation oracle OffDE M = MLE M , the outputs { π t } t ∈ [ T ] of Algorithm 1 satisfies that with probability at least 1 -δ ,

with O ( H log T ) number of oracle calls to the MLE M oracle for δ ∈ (0 , 1 / 2) .

The theorems above show that Algorithm 1 with both epoch schedules achieve near-optimal statistical complexity that a matches the lower bound of Ω( √ HSAT log |M| / log A ) proven by Levy and Mansour [24] up to a poly( H,S,A ) factor.

Computational efficiency. Consider the epoch schedule τ m = 2 m for m ∈ N as discussed in Theorem 3.2. For any unknown T , our algorithm operates over O (log T ) epochs, making one oracle call per epoch. Thus, the computational complexity is O (log T ) oracle calls over T rounds, with an additional per-round cost of O (poly( H,S,A, log T )) . This offers potential advantages over existing algorithms that achieve near-optimal rates without assumptions beyond realizability. The E2D

Algorithm 1 Layerwise pOLicy cover Inverse gaP weighting with trusted OccuPancy measures (LOLIPOP)

Require: epoch schedule 0 = τ -1 = τ 0 < τ 1 < · · · < τ N = T/H , confidence parameter δ ∈ (0 , 1 / 2) , model class M , offline oracle OffDE M .

1: Initialize: ̂ M 0 = { ̂ P h 0 , ̂ R 0 } H h =1 , where ̂ P h 0 is any transtion kernel and ̂ R 0 is constantly 0 .

2: for epoch m = 1 , 2 , · · · , N do

3: Let E m = E M , δ / 2 N 2 ( τ m -1 -τ m -2 ) , γ m = √ H 6 S 4 A 3 / E m and η m = γ m / 720 eH 5 S 3 A 2 .

4: for segment h = 1 , . . . , H do

5: for round t = τ m -1 H +( τ m -τ m -1 )( j -1) + 1 , · · · , τ m -1 H +( τ m -τ m -1 ) h do 6: Observe context c t ∈ C from the environment.

7: for s, a ∈ S × A do 8: Compute

where ˜ d h m is the trusted occupancy measure defined as in Definition 3.1.

9: Let the policy cover Π h m ( c t ) = { ̂ π m -1 ,c t } ∪ { π h,s,a m,c t } s,a ∈S×A .

10: Define p h m ( c t ) to be the Inverse Gap Weighting distribution on the policy cover Π h m,c t

where λ h m,c t is the constant that normalize the distribution. 11: Sample and execute π t ∼ p h m ( c t ) and observe the trajectory c t , π t , { s j t , a j t , r j t } j ∈ [ H ] . 12: j j j

Run OffDE M with the input trajectories { c t , π t , { s t , a t , r t } j ∈ [ H ] } t : m ( t )= m and obtain the h -th layer estimator ̂ P h m and ̂ R h m .

algorithm [16], for instance, requires O ( T ) calls to an online density estimation oracle, involving significantly more calls to a more complex oracle for a general model class M . On the other hand, the Version Space Averaging + E2D algorithm [18] requires O ( T ) calls to an offline density estimation oracle and incurs a computational cost scaling with O ( |M| ) per round. Compared to our algorithm, this results in far more oracle calls and considerably higher computational costs per round.

If the total number of rounds T is known to the learner, we can further reduce the computational cost of LOLIPOP. For any T ∈ N , consider the epoch schedule τ m = 2( T/H ) 1 -2 -m as in Theorem 3.1, similar to Simchi-Levi and Xu [36]. In this scenario, LOLIPOP will run in O (log log T ) epochs, making only O (log log T ) oracle calls over T rounds while still maintaining a slightly worse regret guarantee.

Lower bound on switching cost. There is a lower bound on the switching cost of the scale Ω(log log T ) [44], where the switching cost is the number of switches in the learner's randomized policy. Thus, any learner that only switches its randomized policy after an oracle call will need more than Ω(log log T ) number of oracle calls.

Technical challenge. The main technical challenge is to accurately estimate the occupancy measures for all layers. Naively, the upper bound of divergence between the true occupancy measure and the estimated occupancy measure accumulates exponentially with respect to the number of layers because, naively, the divergence at each layer is upper bounded by the summation of divergences from previous steps. This phenomenon is unavoidable when estimating all the occupancy measures from one dataset generated from a single policy. To avoid this exponential divergence, we apply two methods. First, we turn to layerwise design. Specifically, we generate for occupancy measures at each layer a new dataset from a different policy. This alleviates the exponential accumulation of divergence. Second, we turn to multiplicative guarantees between true occupancy measures and the approximated occupancy measures, i.e., they are equivalent up to a small constant. To achieve this, we construct the trusted occupancy measure (see Definition 3.1), which discards the rarely visited state-action pairs. We then use the trusted occupancy measures to guide exploration.

{

̂

P

h

m

-

1

,

̂

R

h

m

-

1

}

h

∈

[

H

]

Π

1

m

, p

1

m

Π

2

m

, p

2

m

......

̂

P

1

m

,

̂

R

1

m

˜

T

1

m

,

˜

d

2

m

̂

P

2

m

,

̂

R

2

m

˜

T

2

m

,

˜

d

3

m

̂

P

H

m

,

̂

R

H

m

3.2 Detailed Construction and Guarantees of Each Component

In this section, we explain in detail our construction and the guarantees of each component along the dependence graph (Figure 1). We first introduce how the estimators { ̂ P h m -1 , ̂ R h m -1 } h ∈ [ H ] from the previous epoch are used in the new epoch. Then we proceed to introduce how to construct p h m ( c t ) given Π h m ( c t ) . Next, we introduce how are ̂ P h m , ̂ R h m obtained given p h m . Subsequently, we present the definition of the set of trusted transitions ˜ T h m and trusted occupancy measure ˜ d h m . Finally, we present how the policy cover Π h m ( c t ) is constructed.

During epoch m , we will be using the estimators { ̂ P h m -1 , ̂ R h m -1 } H h =1 from the previous epoch for regret estimation. More specifically, for π, c, h, s , we define the value functions with respect to the model { ̂ P h m -1 ( c ) , ̂ R h m -1 ( c ) } H h =1 as ̂ V h m -1 ( s ; π, c ) . The optimal value function is denoted by ̂ V 1 m -1 ( s ; c ) = max π ̂ V 1 m -1 ( s ; π, c ) . For h = 1 , we further simply the notation by denoting ̂ V 1 m -1 ( c ) = ̂ V 1 m -1 ( s 1 ; c ) and ̂ V 1 m -1 ( π, c ) = ̂ V 1 m -1 ( s 1 ; π, c ) . Also denote the optimal policy under context c by ̂ π m -1 ,c = arg max π ̂ V 1 m -1 ( π, c ) . Thus, the regret is estimated to be

At round t , let m ( t ) and h ( t ) be the epoch in which the segment round t lies. We note that during each epoch m and segment h , all of the notions Π h m ( c t ) , ̂ P h m ( c t ) , ̂ R h m ( c t ) , ˜ T h m ( c t ) , ˜ d h m ( · , · ; · , c t ) , p h m ( c t ) , and π h,s,a m,c t will not depend on the specific time step t , but only the context c t . Thus, we will use Π h m ( c ) to denote the policy cover if c t = c when m ( t ) , h ( t ) = m,h . Similar conventions regarding the context c apply to the notations ̂ P h m , ̂ R h m , ˜ T h m , ˜ d h m , p h m , π h,s,a m, · . Under any context c , the policy cover Π h m will include ̂ π m -1 ,c and has no more than SA +1 policies. These two properties together guarantee that the Inverse Gap Weighting [14] randomized policy p h m ( c ) (Line 10) satisfies the following guarantee on the estimated regret.

Lemma 3.1. For any m , h , c , the definition of the randomized policy p h m ( c ) is well defined, i.e., there exist λ h m,c ∈ [0 , 2 SA ] such that ∑ π ∈ Π h m ( c ) p h m,c ( π ) = 1 . Furthermore, we have the estimated regret is bounded by E π ∼ p h m ( c ) [ ̂ reg m -1 ( π, c ) ] ≲ √ H 4 S 4 A 3 · E m .

The choice of λ h m,c here is for compactness of presentation. It can be chosen to be 2 SA for suboptimal arms and collect the probability remained to the optimal arm [36], which is computationally efficient. The computation for the policy π h ( t ) ,s,a m ( t ) ,c t for any t, s, a can be computed in poly( H,S,A, log T ) time by formulating it as a linear fractional programming problem. We defer the details to Appendix G.

Since p h m maps C to randomized policies, it is thus a policy kernel. This means the trajectories generated in epoch m and segment h follow an i.i.d. distribution as described in the definition of Definition 2.1. By applying the guarantee from Definition 2.1, we have the following guarantee on ̂ P h m , ̂ R h m .

Lemma 3.2. For any m , h , and c ∼ D , π ∼ p h m ( c ) , we have with probability at least 1 -δ 2 N 2 , that

If the offline density estimation oracle is chosen to be the Maximum Likelihood Estimation oracle MLE M , we will obtain E m ≲ log( T M /δ ) / ( τ m -1 -τ m -2 ) .

The most involved part of our construction concerns the idea of trusted transitions and trusted occupancy measures . This construction eliminates the parts of transitions that are too scarcely visited. The purpose will be clear in the guarantees (Lemmas 3.3 and 3.4) subsequent to the definitions.

Definition 3.1. For any m,π,c,h,s, a , we define iteratively the trusted occupancy measures ˜ d h m ( s ; π, c ) , ˜ d h m ( s, a ; π, c ) and the set of trusted transitions ˜ T h m ( c ) at layer h as the following:

For any m,h,c , the set of trusted transitions are defined as the set of transtions

where ζ m = γ m 8 eH ( H +1) 2 . Notice that to define ˜ d h m ( s ; π, c ) and ˜ d h m ( s, a ; π, c ) , we only need { ˜ T j m ( c ) , ̂ P j m ( c ) } j ∈ [ h -1] . Thus the two definitions are iteratively well-defined. Meanwhile, we also define the observable occupancy measures as the occupancy measures of the true model going through only the trusted transitions, i.e.,

The computation of the set of trusted transitions need not enumerate all policies. The trusted transition set can be computed in poly( H,S,A, log T ) time by formulating it as a linear fractional programming problem. We defer the details to Appendix G.

Define the estimated occupancy measures ̂ d h m ( s ; π, c ) := E ̂ M m ,π,c [ ✶ ( s h 1 = s )] and ̂ d h m ( s, a ; π, c ) := E ̂ M m ,π,c [ ✶ ( s h 1 , a h 1 = s, a )] . The trusted occupancy measure, though it eliminates rarely visited transitions, remains a valid estimate for all policies because the divergence between the estimated occupancy measure and itself is bounded. Specifically, we have the following lemma.

Lemma 3.3. For all m,π,h,s,a , under any context c , we have

The next guarantee is the key to our analysis and is the most non-trivial guarantee of our construction. The following lemma states that if, for a context c , the Hellinger divergence between ̂ P and P h ⋆ at layer h is small for all h ∈ [ H ] , then the trusted occupancy measure is upper bounded by a scaling of the observable occupancy measure.

Lemma 3.4. For any m and c , assume for all h ∈ [ H ] ,

Then for the same m,c and all π, h, s, a , we have

Since the trusted occupancy measure is upper bounded up to scaling by the observable occupancy measure, then the state-action pairs with large trusted occupancy measures are guaranteed to be visited often in the true dynamics as well. This enables more accurate planning and is thus crucial to our analysis.

Finally, we state the coverage guarantee achieved by the construction of Π h m and p h m . Concretely, we upper bound the trusted occupancy measure ˜ d h m ( · , · ; π, · ) of any policy π by the trusted occupancy measure induced by policy kernel p h m .

Lemma 3.5. For any m,π,c,h,s, a , we have

where E ′ m = ( 2 e 2 √ E m H 4 S 2 A + 1 720 H 4 S 3 A 2 ̂ reg m -1 ( π, c ) ) ˜ d h m ( s, a ; π, c ) ,and p h m ( c, π h,s,a m,c ) is the probability of p h m ( c ) on π h,s,a m,c . The guarantee Eq. (4) also holds replacing D H ( P h ⋆ ( s, a ; c ) , ̂ P h m ( s, a ; c )) with D H ( R h ⋆ ( s, a ; c ) , ̂ R h m ( s, a ; c )) on both sides of the inequality.

4 Regret Analysis

In this section, we prodive a proof sketch of the regret analysis. Detailed proofs are deferred to Appendix E. We first aggregate the component-wise guarantees (Lemmas 3.1 to 3.5) from Section 3 to present the following epoch-wise guarantee.

Lemma 4.1. For any m , any policies { π c } c ∈C , and δ ∈ (0 , 1 / 2) , with probability at least 1 -δ/M ,

Proof sketch of Lemma 4.1. For simplicity, in this proof sketch, we assume the true reward distribution is known 3 . For this, we first apply the celebrated local simulation lemma (Lemma C.2) in reinforcement learning to relate the divergence of the value functions to the stepwise divergences as the following. Under any context c ,

Then we can exchange the estimated occupancy measure ̂ d h m ( s, a ; π c , c ) by the trusted occupancy measure through Lemma 3.3, that is,

Then by coverage guarantee Lemma 3.5, for any h, s, a , we can bound

If the assumption in Lemma 3.4 is satisfied, then by Lemma 3.4 and the definition of p h m , we have

If the assumptions in Lemma 3.4 are not satisfied, we have similar control as well (see the full proof Appendix E for details). Altogether with taking expectation on c , by the offline density estimation bound from Lemma 3.2, we have

A revised version of the regret analysis (Lemma E.1) in Simchi-Levi and Xu [36], which relates the epoch-wise guarantee to the regret estimation error, can be found in Appendix E. Combining Lemmas 4.1 and E.1, we arrive at the following general regret guarantee.

Theorem 4.1. The outputs { π t } t ∈ [ T ] of Algorithm 1 satisfies with probability at least 1 -δ that

for δ ∈ (0 , 1 / 2) .

5 Extension: Reward-free Reinforcement Learning for CMDPs

In this section, we introduce the application of Algorithm 1 in the task of reward-free reinforcement learning in (stochastic) CMDPs. All proofs in this section will be deferred to Appendix F.

Reward-free reinforcement learning [22]. Reward-free reinforcement learning aims to efficiently explore the environment without relying on observed rewards. By doing so, it aims to enable the computation of a nearly optimal policy for any given reward function, utilizing only the trajectory data collected during exploration and without needing further interaction with the environment. This approach holds particular significance in scenarios where reward functions are refined over multiple iterations to encourage specific behaviors through trial and error, such as in constrained RL formulations. In such cases, repeatedly applying the same RL algorithm with varying reward functions can prove to be highly inefficient regarding sample usage, underscoring the efficiency of reward-free reinforcement learning.

Problem formulation. The major differences between the regret minimization setting (Section 2.1) and the reward-free reinforcement learning are that in the latter, no reward signals are observed during the interaction, and the goal of the latter is to output a CMDP prediction ̂ M whose value functions are close to the underlying true CMDP M ⋆ for any reward distributions. To accommodate such a change, for any model M = { P h M , R h M } h ∈ [ H ] and reward distribution R = { R h } h ∈ [ H ] , we define M ( · ; R ) = { P h M , R h } h ∈ [ H ] = { P h M ( c ) , R h ( c ) } h ∈ [ H ] ,c ∈C to be model M with the reward distribution part replaced by R . Thus in the reward-free reinforcement learning setting, the underlying true model satisfies M ⋆ = M ⋆ ( · ; 0) where 0 is used to denote the reward distribution that is constantly 0 .

For any model M reward distribution R , context c and policy π , we use M ( π, c ; R ) to denote the distribution of the trajectory c 1 , π 1 , s 1 1 , a 1 1 , r 1 1 , . . . , s H 1 , a H 1 , r H 1 given M ⋆ = M ( · ; R ) , c 1 = c , and π 1 = π . Also denote the probability and the expectation under M ( π, c ; R ) to be P M,π,c,R ( · ) and E M,π,c,R [ · ] respectively. Given reward distribution R , any policy π , state s and action a , we define the action value function Q h ⋆ ( s, a ; π, c, R ) at layer h and the value function V 1 ⋆ ( s ; π, c, R ) at layer h under context c and policy π as

We denote the optimal policy with reward distribution R under context c as π ⋆,c,R and abbreviate its value function as V h ⋆ ( · ; c, R ) . For h = 1 , we denote V 1 ⋆ ( c, R ) = V 1 ⋆ ( s 1 ; c, R ) and V 1 ⋆ ( π, c, R ) = V 1 ⋆ ( s 1 ; π, c, R ) . We also denote V h M as the value functions when M ⋆ = M .

Assumption 5.1 (Realizability for reward-free RL) . Suppose the learner is given a model class M that contains the underlying true model M ⋆ . Assume all models M ∈ M have 0 reward.

For a given ε, δ > 0 and a model class M , the goal of the learner is to output a model ̂ M at the end of the interaction such that for any reward distribution R and set of policies { π c } c ∈C , the model satisfies

with probability at least 1 -δ . An algorithm that achieves this objective is called ( ε, δ ) -learns the model class M . Then we have the following guarantee from Algorithm 1.

Theorem 5.1. If we choose τ 1 = T/ (2 H ) and τ 2 = T/H , the outputs ̂ M 2 of Algorithm 1 satisfies the reward-free objective Eq. (5) with probability at least 1 -δ , with T at most bounded by

for δ ∈ (0 , 1 / 2) . Moreover, the algorithm requires O ( H ) number of oracle calls to the MLE M oracle.

The proof follows a similar argument of Lemma 4.1. In addition, we have a matching lower bound up to a poly( H,S,A ) factor adapted from the non-contextual lower bound from Jin et al. [22].

Theorem 5.2. Fix ε ≤ 1 , δ ≤ 1 / 2 , H,A ≥ 2 . Suppose S ≥ L log A for a large enough universal constant L and K ≥ 0 large enough. Then, there exists a CMDP class M with |S| = S , |A| = A , |M| = K , and horizon H and a distribution µ on M such that any algorithm Alg that ( ε/ 24 , δ ) -learns the class M satisfies E M ∼ µ E M, Alg [ T ] ≳ log |M| /ε 2 , where T is the number of trajectories required by the algorithm Alg to achieve ( ε/ 24 , δ ) accuracy and E M, Alg [ · ] is the expectation under the interaction between the algorithm Alg and model M .

6 Discussion

Extension to low rank CMDPs. Low rank MDPs represent a significant extension to tabular MDPs, as explored in various studies [6, 28, 21, 4]. Linear MDPs are typically the first step beyond tabular MDPs. Extending our approach to linear CMDPs would be a substantial achievement. The primary challenge lies in identifying the trusted transitions within linear CMDPs. The current construction for tabular CMDPs does not readily apply here because it does not utilize the low-rank structure.

Extension to model-free learning. Our approach is model-based. However, model-free methods are often more practical for real-world applications. The main challenge lies in effectively balancing exploration and exploitation using only the value functions, as opposed to our method which depends on the occupancy measure.

More efficient oracles. In this paper, we focus on offline density estimation oracles due to the necessity of a small Hellinger distance between the estimated model and the true model for our approach. An offline regression oracle would only provide a 2-norm distance guarantee, which is inadequate for our purposes. Nevertheless, it is interesting to explore whether a reduction from CMDPs to offline regression could be feasible.

References

Alekh Agarwal. Selective sampling algorithms for cost-sensitive multiclass prediction. In International Conference on Machine Learning , pages 1220-1228. PMLR, 2013.

Alekh Agarwal, Miroslav Dudík, Satyen Kale, John Langford, and Robert E Schapire. Contextual bandit learning with predictable rewards. In Artificial Intelligence and Statistics , 2012.

Alekh Agarwal, Daniel Hsu, Satyen Kale, John Langford, Lihong Li, and Robert Schapire. Taming the monster: A fast and simple algorithm for contextual bandits. In International Conference on Machine Learning , pages 1638-1646, 2014.

Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. FLAMBE: Structural complexity and representation learning of low rank MDPs. Neural Information Processing Systems (NeurIPS) , 2020.

Philip Amortila, Dylan J. Foster, and Akshay Krishnamurthy. Scalable online exploration via coverability, 2024.

Steven J Bradtke and Andrew G Barto. Linear least-squares algorithms for temporal difference learning. Machine learning , 22(1):33-57, 1996.

Abraham Charnes and William W Cooper. Programming with linear fractional functionals. Naval Research logistics quarterly , 9(3-4):181-186, 1962.

Fan Chen, Song Mei, and Yu Bai. Unified algorithms for rl with decision-estimation coefficients: No-regret, pac, and reward-free learning. arXiv preprint arXiv:2209.11745 , 2022.

Jinglin Chen, Aditya Modi, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. On the statistical efficiency of reward-free exploration in non-linear rl. Advances in Neural Information Processing Systems , 35:20960-20973, 2022.

Yuan Cheng, Ruiquan Huang, Jing Yang, and Yingbin Liang. Improved sample complexity for reward-free reinforcement learning under low-rank mdps. arXiv preprint arXiv:2303.10859 , 2023.

Junze Deng, Yuan Cheng, Shaofeng Zou, and Yingbin Liang. Sample complexity characterization for linear contextual mdps. arXiv preprint arXiv:2402.02700 , 2024.

Simon S Du, Yuping Luo, Ruosong Wang, and Hanrui Zhang. Provably efficient Q-learning with function approximation via distribution shift error checking oracle. In Advances in Neural Information Processing Systems , pages 8060-8070, 2019.

Miroslav Dudik, Daniel Hsu, Satyen Kale, Nikos Karampatziakis, John Langford, Lev Reyzin, and Tong Zhang. Efficient optimal learning for contextual bandits. In Proceedings of the Twenty-Seventh Conference on Uncertainty in Artificial Intelligence , pages 169-178. AUAI Press, 2011.

Dylan J Foster and Alexander Rakhlin. Beyond UCB: Optimal and efficient contextual bandits with regression oracles. International Conference on Machine Learning (ICML) , 2020.

Dylan J Foster, Alekh Agarwal, Miroslav Dudík, Haipeng Luo, and Robert E. Schapire. Practical contextual bandits with regression oracles. International Conference on Machine Learning , 2018.

Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487 , 2021.

Dylan J Foster, Alexander Rakhlin, Ayush Sekhari, and Karthik Sridharan. On the complexity of adversarial decision making. Advances in Neural Information Processing Systems , 35: 35404-35417, 2022.

Dylan J Foster, Yanjun Han, Jian Qian, and Alexander Rakhlin. Online estimation via offline estimation: An information-theoretic framework. arXiv preprint arXiv:2404.10122 , 2024.

Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes, 2015.

Pihe Hu, Yu Chen, and Longbo Huang. Towards minimax optimal reward-free reinforcement learning in linear mdps. In The Eleventh International Conference on Learning Representations , 2022.

Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contextual decision processes with low Bellman rank are PAC-learnable. In International Conference on Machine Learning , pages 1704-1713, 2017.

Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. Reward-free exploration for reinforcement learning. In International Conference on Machine Learning , pages 4870-4879. PMLR, 2020.

Yin Tat Lee and Aaron Sidford. Solving linear programs with sqrt(rank) linear system solves. arXiv preprint arXiv:1910.08033 , 2019.

Orin Levy and Yishay Mansour. Optimism in face of a context: Regret guarantees for stochastic contextual mdp, 2023.

Orin Levy, Asaf Cassel, Alon Cohen, and Yishay Mansour. Eluder-based regret for stochastic contextual mdps. arXiv preprint arXiv:2211.14932 , 2022.

Orin Levy, Alon Cohen, Asaf Cassel, and Yishay Mansour. Efficient rate optimal regret for adversarial contextual mdps using online function approximation. In International Conference on Machine Learning , pages 19287-19314. PMLR, 2023.

Gen Li, Yuling Yan, Yuxin Chen, and Jianqing Fan. Optimal reward-agnostic exploration in reinforcement learning. 2023.

Francisco S Melo and M Isabel Ribeiro. Q-learning with linear function approximation. In International Conference on Computational Learning Theory , pages 308-322. Springer, 2007.

Zakaria Mhammedi, Dylan J Foster, and Alexander Rakhlin. Representation learning with multi-step inverse kinematics: An efficient and optimal approach to rich-observation rl. In International Conference on Machine Learning , pages 24659-24700. PMLR, 2023.

Zakaria Mhammedi, Adam Block, Dylan J. Foster, and Alexander Rakhlin. Efficient model-free exploration in low-rank mdps, 2024.

Sobhan Miryoosefi and Chi Jin. A simple reward-free approach to constrained reinforcement learning. In International Conference on Machine Learning , pages 15666-15698. PMLR, 2022.

Aditya Modi and Ambuj Tewari. Contextual markov decision processes using generalized linear models. arXiv preprint arXiv:1903.06187 , 2019.

Aditya Modi, Nan Jiang, Satinder Singh, and Ambuj Tewari. Markov decision processes with continuous side information. In Algorithmic Learning Theory , pages 597-618. PMLR, 2018.

Yury Polyanskiy and Yihong Wu. Lecture notes on information theory. Lecture Notes for ECE563 (UIUC) and , 6(2012-2016):7, 2014.

Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming . John Wiley & Sons, 2014.

David Simchi-Levi and Yunzong Xu. Bypassing the monster: A faster and simpler optimal algorithm for contextual bandits under realizability. Mathematics of Operations Research , 2021.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018.

Andrew Wagenmaker, Max Simchowitz, and Kevin Jamieson. Beyond no regret: Instancedependent pac reinforcement learning, 2022.

Ruosong Wang, Simon S Du, Lin Yang, and Russ R Salakhutdinov. On reward-free reinforcement learning with linear function approximation. Advances in neural information processing systems , 33:17816-17826, 2020.

Yunbei Xu and Assaf Zeevi. Upper counterfactual confidence bounds: a new optimism principle for contextual bandits. arXiv preprint arXiv:2007.07876 , 2020.

Yunbei Xu and Assaf Zeevi. Upper counterfactual confidence bounds: a new optimism principle for contextual bandits, 2024.

Xuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang, Alekh Agarwal, and Wen Sun. Efficient reinforcement learning in block mdps: A model-free representation learning approach. In International Conference on Machine Learning , pages 26517-26547. PMLR, 2022.

Zihan Zhang, Simon S Du, and Xiangyang Ji. Nearly minimax optimal reward-free reinforcement learning. arXiv preprint arXiv:2010.05901 , 2020.

Zihan Zhang, Yuhang Jiang, Yuan Zhou, and Xiangyang Ji. Near-optimal regret bounds for multi-batch reinforcement learning. Advances in Neural Information Processing Systems , 35: 24586-24596, 2022.

A Notation

For any integer n , we use [ n ] to denote the set { 1 , . . . , n } . For any set X , we use ∆( X ) to denote the set of all distributions on the set X . We define O ( · ) , Ω( · ) , Θ( · ) , ˜ O ( · ) , ˜ Ω( · ) , ˜ Θ( · ) following standard non-asymptotic big-oh notation. We use the binary relation x ≲ y to indicate that x ≤ O ( y ) . ✶ ( E ) is an indicator function of event E .

B Related Works

Contextual bandits and contextual MDPs. The SquareCB [14] obtains optimal regret for contextual bandits with access to an online regression oracle. This is extended to the CMDPs by Foster et al. [16] with the E2D algorithm. However, the algorithm requires O ( T ) called to an online density estimation oracle. Compared to our algorithm, it necessitates significantly more calls to an oracle that is harder to implement for a general model class M . After the FALCON algorithm [36] establishes the reduction from contextual bandits to offline regression, Xu and Zeevi [41] proposed the UCCB algorithm, which is less oracle-efficient in terms of oracle calls but adopts the prevalent "optimism in the face of uncertainty" principle and is thus easier to generalize. More specifically, the UCCB algorithm is extended to CMDPs by Levy and Mansour [24], Deng et al. [11] with assumptions on the model class. The RM-UCDD algorithm proposed by Levy and Mansour [24] requires the model class to have a minimum reachability p min to all states under any policy and the regret guarantees scale with O (poly( H,S,A ) · (1 /p min ) √ T log |M| ) . This assumption precludes model classes with small reachability, which frequently happens in practice [4]. The CMDP-VR algorithm proposed by Deng et al. [11] assumes a varying representation assumption on the model class instead. The assumption asserts that any model M = { P h , R h } h ∈ [ H ] ∈ M satisfies for any c, h, s, a, s ′ , P h ( s ′ | s, a ; c ) = ⟨ ϕ h ( s, a ; c ) , µ h ( s ′ ) ⟩ for the known feature vector ϕ h ( s, a ; c ) ∈ R d and an unknown vector µ h ( s ′ ) ∈ R d which does not depend on the context c . This assumption is stringent because canonically, the feature vector for CMDPs will be chosen to be the unit vector indexed by s, a in R SA , i.e., ϕ h ( s, a ; c ) = e s,a ∈ R SA . Then, the requirement of µ h ( s ′ ) not depending on the context forces P h to not depend on the context as well. This reduces the CMDP to an MDP. While it is possible to complicate the feature vector to not reduce to an MDP, this would result in a higher dimension in the feature vector space, which will be reflected in the regret bounds obtained ( ˜ O (poly( H,d ) √ T log |M| ) . Another significant disadvantage compared with our algorithm is that the RM-UCDD and CMDP-VR algorithm requires O ( T ) number of oracle calls. Other structural assumptions on the model class, such as a generalized linear structure or bounded eluder dimension, have been explored by Modi and Tewari [32], Levy et al. [25].

Reward-free reinforcement learning. Reward-free reinforcement learning aims to efficiently explore the environment without relying on observed rewards. By doing so, it aims to enable the computation of a nearly optimal policy for any given reward function, utilizing only the trajectory data collected during exploration and without needing further interaction with the environment. This framework is proposed by [22] and has been extensively studied for MDPs with various assumptions [43, 4, 39, 42, 9, 31, 38, 20, 10, 27, 30, 5]. However, to the best of our knowledge, we are the first to study the reward-free reinforcement learning setting for stochastic CMDPs. We provide a near-optimal sample complexity upper bound and a matching lower bound up to a poly( H,S,A ) factor with only O ( H ) number of oracle calls. Nevertheless, the upper bound is obtained by adjusting the exploration-exploitation, highlighting the flexibility of our algorithm.

C Technical Tools

C.1 Maximum Likelihood Estimation for Density Estimation

Example C.1 (MLE for finite model class) . Let M be a finite model class and the MLE estimator ̂ M be defined by

For any δ ∈ (0 , 1 / 2) , we have with probability at least 1 -δ ,

Proof of Example C.1. For any δ ∈ (0 , 1 / 2) , by Lemma C.2 of Foster et al. [18], we have with probability at least 1 -δ/ 2 ,

Then by Lemma A.3 of Foster et al. [16], we have with probability at least 1 -δ/ 2 ,

Then by union bound, we obtain the desired result.

C.2 Information Theory

Lemma C.1 (Lemma B.4 of Foster et al. [17]) . Let P and Q be two distributions on space χ . Let h : χ → R be a function. Then we have:

C.3 Reinforcement Learning

Lemma C.2 (Lemma F.3 of [16]) . Let M = { P h M , R h M } h ∈ [ H ] and M = { P h M , R h M } h ∈ [ H ] be two CMDPs. For any policy π and context c , we have

D Proofs from Section 3

In this section, we present the proofs for Lemmas 3.1 to 3.5.

Proof of Lemma 3.1. We fix an arbitrary context c throughout the proof. Let u ( λ ) := ∑ π ∈ Π h m 1 / ( λ + η m · ̂ reg m -1 ( π )) . Since for any h ∈ [ H ] , ̂ π m -1 ,c ∈ Π h m , then u ( λ ) ≥ 1 / ( λ + η m · ̂ reg m -1 ( ̂ π m -1 ,c )) = 1 /λ . On the other hand, u ( λ ) ≤ ( SA + 1) /λ . Moreover, u ( λ ) is clearly monotonically decreasing with u (0) = ∞ and u ( SA +1) ≤ 1 . Thus there exists λ h m,c ∈ (0 , SA +1] such that u ( λ h m,c ) = 1 as we desired.

Now we have the regret is bounded by

Finally, we recall η m = γ m / (720 e 3 H 5 S 3 A 2 ) and γ m = √ H 6 S 4 A 3 E m plug this bound in, then we have

Proof of Lemma 3.2. By definition of Definition 2.1.

Proof of Lemma 3.3. For any m,π,c,h,s, a , by the definition of ˜ d h m ( s, a ; π, c ) , we the difference between ̂ d h m ( s, a ; π, c ) and ˜ d h m ( s, a ; π, c ) are the parts of occupancy measures that do not go through the trusted transitions, i.e.,

where ̂ P j +1: h m +1 ( s | s j +1 ; π, c ) is the estimated transition probability from s j +1 at step j +1 to s at step h under policy π and context c . Then since ( s j , a j , s j +1 ) / ∈ ˜ T j m ( c ) , we have

Proof of Lemma 3.4. We prove iteratively between the objective

and the following claim: For any h and any ( s, a, s ′ ) ∈ ˜ T h m ( c ) , we have

First, we show that for any h ∈ [ H ] ,

For this, we note by Lemma C.1 that for any h, s, a, s ′ ,

where the second inequality is by data-processing inequality [34]. Then by AM-GM, we have

Plug the above back into Eq. (8) and reorganize, we obtain

Then multiply both sides by ˜ d h m ( s, a ; π, c ) , we have

Meanwhile, by the definition of π h,s,a m,c , we have

Then by the induction hypothesis, we have

h

where the last equality is by the definition of ζ m . In all, we have shown that

Now we prove by contradiction, if for any ( s, a, s ′ ) ∈ ˜ T h m ( c ) , the reverse inequality is true, i.e., ̂ P h m ( s ′ | s, a ; c ) > ( 1 + 1 H ) 2 P h ⋆ ( s ′ | s, a ; c ) . Then for any π ,

This contradicts the definition of ˜ T h m ( c ) .

For the other direction of Eq. (6), we prove for all h ∈ [ H ] ,

This direction is straightforward since

With the two derivations Eq. (6) and Eq. (9), along with the fact that the initial argument of Eq. (6) holds by defintion for h = 1 . Thus we conclude the proof.

Proof of Lemma 3.5. For any m,π,c,h,s, a , by AM-GM, we have

By the definition of π h,s,a m,c , we have futher

Recall the choice of η m = γ m / (720 e 3 H 5 S 3 A 2 ) and γ m = √ H 6 S 4 A 3 E m , we have

Also by the definition of p h m ( c ) , we have

Now we plug Eqs. (11) to (13) back into Eq. (10) to obtain that

Similar bounds can be obtained replacing D H P h ⋆ ( s, a ; c ) , ̂ P h m ( s, a ; c ) with D H ( R h ⋆ ( s, a ; c ) , ̂ R h m ( s, a ; c ) ) .

E Proofs from Section 4

Proof of Lemma 4.1. For this we first apply the local simulation lemma (Lemma C.2) in reinforcement learning to relate the divergence of the value functions to the stepwise divergences as the following. Under any context c ,

Then we can exchange the estimated occupancy measure ̂ d h m ( s, a ; π c , c ) by the trusted occupancy measure through Lemma 3.3, that is,

Then by coverage guarantee Lemma 3.5, for any h, s, a , we can bound

Suppose the assumption in Lemma 3.4 is satisfied, then by Lemma 3.4, we have

The D H ( ̂ P h m ( s h 1 , a h 1 ; c ) , P h ⋆ ( s h 1 , a h 1 ; c ) ) in Eqs. (15), (16) and (26) can be replaced with D H ( ̂ R h m ( s h 1 , a h 1 ; c ) , R h ⋆ ( s h 1 , a h 1 ; c ) ) as well. If the assumption in Lemma 3.4 is not satisfied, then we there exists j such that

This implies

Thus, altogether, no matter the assumption in Lemma 3.4 is satisfied or not, we have shown

Taking expectation on c , we have

where the second inequality is by the offline density estimation bound from Lemma 3.2

Lemma E.1 ([36]) . Let ε 1 , . . . , ε N be N positive values. Suppose for any m> 0 and an arbitrary policy set { π c } c ∈C , we have:

Then for any m> 0 ,

where δ 1 = 2 ε 1 + 1 10 and δ m = 1 9 δ m -1 + 20 9 ε m for any m ≥ 2 .

Proof of Lemma E.1. We present the proof here for completeness. By Eq. (17), we have that for all m ≥ 0 and π c ,

Symmetrically, we have

Now we inductively show for m = 1 , 2 , . . . that Eq. (18) and Eq. (19) hold. For m = 1 , since E c ∼D [reg( π c , c )] , E c ∼D [ ̂ reg 1 ( π c , c )] ≤ 1 for all π , then we have from Eq. (20) and Eq. (21)

Hence we have shown Eq. (18) and Eq. (19) for m = 1 and δ 1 = 2 ε 1 + 1 10 . Now, suppose Eq. (18) and Eq. (19) holds for all 1 , 2 , . . . , m -1 . Plugging Eq. (19) for m -1 into the right hand side of Eq. (20), we have

where the last equality is by E c ∼D [reg( π ⋆,c , c )] = 0 . Thus reorganizing the terms, we have

where δ m = 1 9 δ m -1 + 20 9 ε m . Thus we have shown Eq. (18) for m . Then plugging Eq. (19) for m -1 into the right hand side of Eq. (21), we have

Furthermore, by Eq. (18) for m we have E c ∼D [reg( ̂ π m,c , c )] ≤ 10 9 E c ∼D [ ̂ reg( ̂ π m,c , c )] + δ m = δ m . Plug this in the aforementioned inequality, we have

Reorganizing the terms, this in turn gives

where recall δ m = 1 9 δ m -1 + 20 9 ε m . This proves Eq. (19) for m which completes our induction.

Proof of Theorem 4.1. Let E 0 = 1 . Then we have by Lemma 4.1, Eq. (17) holds with ε m = L √ H 6 S 4 A 3 E m for L > 0 large enough for all m> 0 . Combining Lemmas 3.1 and E.1, we have

Then by Lemma 3.1, we have further

Hence, we have

In all, we can obtain the following regret bound with probability at least 1 -δ ,

where the last step takes a union bound on the offline oracle guarantees.

Proof of Theorem 3.1. Without loss of generality, assume that T/H > 1000 . Since we are choosing τ m = 2( T/H ) 1 -2 -m , we first note that since τ m ≤ T/H , we have ( T/H ) 2 -m ≥ 2 . Then we have

This implies τ m -τ m -1 ≥ 1 2 τ m for m ∈ [ N ] . Then by Theorem 4.1, we have with probability at least 1 -δ ,

where the last inequality follows from

and N = O (log log T ) . So the number of oracle calls are O ( H log log T ) .

Proof of Theorem 3.2. Since we are choosing τ m = 2 m , we have N = O (log T ) . So the number of oracle calls are O ( H log T ) . Meanwhile, by Theorem 4.1, we have with probability at least 1 -δ ,

F Proofs from Section 5

Proof of Theorem 5.1. The reward distributions ̂ R m are 0 for m = 0 , 1 since ̂ R 0 are set to be 0 and the model class M consists of models with constantly 0 reward. The regret estimations ̂ reg m -1 are all 0 for m = 1 , 2 . Thus we apply the component-wise guarantees (Lemmas 3.2 to 3.5) to ̂ P 2 , ̂ R 2 where ̂ R 2 = 0 . Concretely, from Lemma 3.2 we have with probability at least 1 -δ ,

with the offline MLE M oracle. From Lemma 3.3, we have for any π, c, h, s, a

,

From Lemma 3.4, we have if for a context c and all h ∈ [ H ] ,

Then for the same c and all π, h, s, a , we have

Finally, from Lemma 3.5, we have for any π, c, h, s, a ,

Now we are ready to prove our claim following similar derivations from the proof of Lemma 4.1. Concretely, we have for any set of policies { π c } c ∈C and any context c , by local simulation lemma (Lemma C.2)

Then by Eq. (23), we have

Case I: If for a context c and all h ∈ [ H ] ,

Then by Eqs. (24) and (25),

Thus we have

Case II : If for a context c there exists j ∈ [ H ] such that

Then

Then take expectation with respect to c together with Eq. (22) we obtain

where the second inequality is by E 2 = H log( |M| /δ ) /T and the last inequality holds when

Thus the reward-free objective Eq. (5) is satisfied with probability at least 1 -δ with

Lemma F.1 (Lemma D.2 of Jin et al. [22]) . Fix ε ≤ 1 , δ ≤ 1 / 2 , H,A ≥ 2 , and suppose that S ≥ L log A for a universal constant L . Then consider the trivial context space C = { c 0 } , there exists a model class M = { M j } j ∈J , with |S| = S , |A| = A , |J | ≤ e SA , and horizon H and a distribution µ on M such that any algorithm Alg that ( ε/ 12 , δ ) -learns the class M satisfies

where T is the number of trajectories required by the algorithm Alg to achieve ( ε/ 12 , δ ) accuracy and E M, Alg [ · ] is the expectation under the interaction between the algorithm Alg and model M .

Proof of Theorem 5.2. . Let n = log K/ ( SA ) , then we consider the context space C = { c 1 , ..., c n } and the i.i.d. distribution D on the context space being uniform. Denote the model class obtained from Lemma F.1 by M = { M j ( c 0 ) } j ∈J . Let J = { j 1 , ..., j n } ∈ J n be an index. Then we consider the model class M = { M J } J ∈J n , where M J ( c i ) = M j i ( c 0 ) , that is, the model class M is on each context c i ∈ C an independent M . We first have that the size of the model class is bounded by |M| = |J | n ≤ e nSA ≤ K . Then we have for any algorithm Alg that ( ε/ 24 , δ ) -learns the class M , it must have ( ε/ 12 , δ ) -learns the class M ( c i ) = { M J ( c i ) } J ∈J n = { M j ( c 0 ) } j ∈J for at least half of the contexts c i ∈ C by Markov's inequality. This, in turn, combined with Lemma F.1 gives that there exists a distribution µ on the model class M such that

where T is the number of trajectories required by the algorithm Alg to achieve ( ε/ 12 , δ ) accuracy and E M, Alg [ · ] is the expectation under the interaction between the algorithm Alg and model M . Thus concludes our proof.

G Computation

In this section, we show that for any m,c,h,s, a , the policy π h,s,a m,c (Line 8) and the trusted transitions ˜ T h m ( c ) (Definition 3.1) can be computed efficiently through linear programming. For simplicity, we fix a context c throughout this section and omit its dependence. We first note that if ˜ T j m for j ≤ h -1 and π h,s,a m are computed, then ˜ T h m can be computed in poly( HSA ) time. To see this, for any ( s, a, s ′ ) , we have by the definition of π h,s,a m that

Then ( s, a, s ′ ) ∈ ˜ T h m iff

where the left-hand side can be computed given ˜ T j m for j ≤ h -1 in poly( HSA ) time. Thus we only need to show how to compute π h,s,a m . Assume we want to compute π ¯ h, ¯ s, ¯ a m for ¯ h ∈ [ H ] , ¯ s ∈ S , ¯ a ∈ A , given ˜ T h m for h ≤ ¯ h -1 . Let ¯ r h m ( s, a ) = E r h ∼ ̂ R h m -1 [ r h ] be the mean reward for model ̂ M m -1 for h, s, a . Let ¯ r m = (¯ r h m ( s, a )) h,s,a be the vector of mean rewards for the model ̂ M m -1 . We consider the following linear fractional program with the following decision variables:

We will use the variable ˜ d ¯ h m to simulate the trusted occupancy measures and ̂ d m -1 to simulate the estimated occupancy measures from ̂ M m -1 by linear constraints. Concretely, consider the following

linear fractional program:

This is a linear fractional program of HSA + ¯ hSA decision variables with HSA + HS + ¯ hSA + ¯ hS constraints. It is clear from the linear constraints that this program simulates the MDPs ̂ M m -1 and ̂ M m , and the program obtains the right objective. Then, for this linear fractional program, we apply the Charnes-Cooper transformation [7] to transform it into a linear program. After the transformation, we can apply existing tools for solving linear programs (e.g., Lee and Sidford [23]) to solve for ˜ d ¯ h m which encodes the policy π ¯ h, ¯ s, ¯ a m in poly( HSA ) time.

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

1. Claims

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

Answer: [Yes] .

Justification: The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope. The claims are validated by detailed proofs.

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: The paper discusses the limitations of the work in the discussion section.

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: The paper provides detailed assumptions and 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 crossreferenced.

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: This 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: This paper does not include experiments.

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: This 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: This 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: This 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, in every respect, with the NeurIPS Code of Ethics.

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 theoretical work. There is no societal impact on the work performed.

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: The paper 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: This 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: This 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: This paper does not involve crowdsourcing or 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: This paper does not involve crowdsourcing or 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:\nIn this paper, the authors study the problem of learning stochastic contextual MDPs using a realizable and finite models class $\\mathcal{M}$ accessed via an offline density estimation oracle, which have the guarantee of minimizing the expected error measured by the squared Hellinger distance. Under the assumption that the oracle is efficient, the presented algorithm is efficient in terms of both running time and oracle call complexity. It achieves a regret bound of  $\\widetilde{O}(\\sqrt{H^7 S^4 A^3 T \\log (|\\mathcal{M}|/\\delta)})$ where $T$ is the number of episodes, $H$ is the horizon, $S,A$ are the finite cardinalities of the state and action spaces, respectively.\nIn this bound, the dependency in $T, |\\mathcal{M}|$ is optimal while the dependency in $H,S,A$ is non-optimal.\nTheir approach is a direct extension of the Inverse-Gap-Weighting (IGW) [Foster and Rakhlin (2020), Simchi-Levi and Xu (2021)] to CMDPs, where using the value function and policy cover instead of the immediate reward to define the approximated suboptimality gap. For the model approximation, they use an offline density estimation oracle (rather than a regression oracle used in precious works on CMAB). \n \nThe authors also show their method applies to the reward-free setting. This is possible as their approch applies strong exploration, which costs in high dependency in $H,S,A$.\n\nSTRENGTHS:\n1.\tThis paper mainly discusses the question of deriving regret bounds for stochastic CMDPs using offline oracles which is an interesting open question in RL theory. \n2.\tThe proposed algorithm successfully extends IGW technique to derive rate-optimal regret bound for CMDPs. Due to the high exploration property of their method, the authors were able to show it also has well performance for the problem of rewards-free exploration.\n\nWEAKNESSES:\n1.\tIn my opinion, the writing requires some improvements. For instance, the algorithm is hard to understand due to some definitions that appear after it (instead of before it). An intuitive explanation of the use in the trusted occupancy measures set is missing, the use of them to derive a multiplicative guarantee is not trivial and not well explained. \n\n2.\tThe authors focus on their logarithmic oracle complexity as their main novelty, whereas in CMDPs, as previous literature shows, oracle complexity of $O(T)$ is considered efficient and acceptable, and the focus should be on using the minimal assumptions regarding the oracle and function class. \n\n3.\tIn continuing to point 2, in this work the authors assume the oracle has a convergence guarantee w.r.t the expected squared Hellinger distance applied on CMDPs class directly, and in Appendix C.1 they provide an implementation using maximum likelihood estimation, which can be hard to optimize for many function classes. Also, throughout the paper, the oracle efficiency is not discussed at all, and no examples of implementations of it for specific easy classes such as linear functions, tabular MDPs, and more are not given.  My main concern is that this oracle is not practical, where the truly interesting question in stochastic CMDPs is to derive regret bound using oracles that potentially can be implementable and efficient (at least for simple classes such as linear functions). \nMoreover, the squared Hellinger distance has a strong relation to the log loss (as established by Foster et al. 2021), and any reference or use in that relation does not appear in this paper. Instead, the authors use an oracle that its efficiency is unclear and not discussed.\nI would appreciate it if the authors would provide an explanation as to why they chose to use a convergence guarantee w.r.t the squared Hellinger distance rather than other convex loss functions (such as the log loss for instance), that clearly yields an efficient oracle for simple function classes such as linear functions. \n\n4. Regarding the computational complexity - in each round the authors compute policy cover for each layer, state, and action. This has a significant cost that is completely ignored.  I would expect that the authors will present their algorithms running time complexity, excluding the oracle's complexity.\n\n5.\tThere are recent, relevant works on CMDPs that are not mentioned in the related literature review. See [1,2,3,4] for references. I also think that a comparison with the results of [2,3,4] might be relevant.\n\n[1] Contextual Markov Decision Processes. Hallak et al. (2015).\n\n[2] Contextual Markov Decision Processes using Generalized Linear Models. Modi and Tewari. (2019).\n\n[3] Efficient rate optimal regret for adversarial contextual MDPs using online function approximation. Levy et al. (2023).\n\n[4] Eluder-based Regret for Stochastic Contextual MDPs. Levy et al. (2024).\n\nQUESTIONS:\n1.\tIn line 4 of Algorithm 1, the segment loop is unclear to me. Can you please explain what is used for? As I understand it, the segment loop is running over the CMDPs layers and hence should be inside the rounds loop. \n2.\tThe oracle guarantee in Definition 2.1 is w.r.t the predicted model $\\hat{M}$ and the true model $M_\\star$. It is unclear to me how it directly implies Lemma 3.2, unless the oracle is actually more powerful than described in definition 2.1 and assumed to optimize simultaneously the squared Hellinger distance of both a rewards function class and a dynamics function class. \nIf this is indeed the situation, it is unclear to me, as the rewards do not define a distribution, and the squared Hellinger distance is defined for distributions. \nAlso, in this case, it is unclear to me why the squared Hellinger distance is an appropriate loss for the rewards approximation. Previous works on CMAB (e.g., Simchi-Levi and Xu 2021, Foster and Rakhlin 2020) show that the least squares (that are less sensitive than the squared Hellinger distance in [0,1]) is an appropriate loss choice for the rewards approximation.  \n 3.\tPlease refer to the concerns regarding the oracle raised in weakness, points 2 and 4. Specifically, I would appreciate an explanation of why you chose this specific oracle and whether it can be implemented efficiently.\n4. I do not understand how the authors derived the fully multiplicative lower bound over the trusted occupancy measures. As I understand it, there is a missing additive term. I would appreciate it if the authors could explain that point as if it does not hold true, the whole analysis is incorrect.\n\nLIMITATIONS:\n1.\tThe regret bound is not tight in $H,S,A$.\n2.\tThe oracle might be inefficient.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-17"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies low-rank contextual decision processes (CMDPs) in offline settings. The authors introduce a novel algorithm that leverages the structure of low-rank CMDPs to achieve efficient learning with limited data. The proposed method, O-LRCDP, is designed to minimize the dependence on large-scale data by utilizing an oracle to provide efficient estimations. The paper presents theoretical guarantees for the performance of O-LRCDP and demonstrates its effectiveness through rigorous analysis and experimental results.\n\nSTRENGTHS:\n1. The paper introduces a novel algorithm (LOLIPOP) for efficiently learning optimal policies in contextual Markov decision processes (CMDPs) with near-optimal regret guarantees.\n\n2. The rigorous theoretical analysis provides strong regret bounds of $O(\\sqrt{H^7S^4A^3T \\log{|M|}})$ that match lower bounds up to polynomial factors in H, S, and A. The algorithm achieves computational efficiency by requiring only O(H log T) or O(H log log T) calls to an offline density estimation oracle, which is a significant improvement over prior work requiring O(T) oracle calls.\n\n3. The paper provides a clear comparison to current algorithms and highlights the key advantages of LOLIPOP over existing ones (Table 1).\nThe algorithm and analysis are versatile, extending beyond regret minimization to the reward-free reinforcement learning setting with near-optimal sample complexity.\n\n4. This paper has a well-organized structure. It adequately discusses relevant work and provides an informative introduction to the basics. I find understanding the main parts of this work quite smooth.\n\n5. The technical approach introducing \"trusted occupancy measures\" to address the multi-layer structure of CMDPs is novel and insightful.\n\nWEAKNESSES:\n1. While the theoretical guarantees are strong, experimental results demonstrating the practical performance and computational efficiency would strengthen the paper significantly, even in a synthetic setting.\n\n2. The discussion of the algorithm's potential limitations or failure cases is limited.\n\n3. The paper uses a standard realizability assumption (Assumption 2.1), which is common in theoretical RL work. While this assumption is reasonable for analysis, a brief discussion of its practical implications and potential ways to verify or approximately satisfy it in real-world scenarios could enhance the paper's applicability.\n\n4. While the extension to reward-free RL is interesting, this section feels somewhat disconnected from the main focus of the paper. The motivation, connection to previous sections, and practical implications of this part could be elaborated further.\n\n5. The presentation of technical details, particularly in Section 3.2, is quite dense and may impede overall readability. While the mathematical rigor is appreciated, the main text could benefit from focusing more on intuitive ideas and key insights that directly address the technical challenges. Consider moving detailed lemmas and heavy mathematical derivations to the appendix. To enhance accessibility and understanding, it would be beneficial to incorporate more high-level explanations and possibly graphical illustrations in the main body.\n\nQUESTIONS:\n1. Have the authors had any chance to perform any empirical investigations comparing LOLIPOP to existing algorithms like E2D or CMDP-VR (even in a synthetic setting)? If so, what were the key findings?\n\n2. The algorithm relies on an offline density estimation oracle. How sensitive is the performance to the choice of this oracle? Are there particular implementations you would recommend in practice?\n\n3. As for the regret bound, are there particular aspects of these dependencies that you think could potentially be improved, or do you believe they are essentially tight given current techniques? Are there specific challenges in tightening the bounds with respect to the state space size S or the model class size |M|?\n\n4. How well do you expect LOLIPOP to perform in settings where the realizability assumption is violated? Are there natural extensions to handle model misspecification?\n\n5. The extension to the reward-free RL setting is interesting. Do the authors see potential applications or connections to other decision-making problems as well (e.g., constrained RL)?\n\nLIMITATIONS:\nThe authors adequately address the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the stochastic Contextual MDP problem under an offline function approximation oracle. Concretely, the authors assume access to an offline density estimation oracle for a (realizable) class of CMDPs. Under this (minimal) assumption, they prove a rate-optimal regret bound, I.e., that scales with $\\sqrt{T \\log |\\mathcal{M}|}$ where $T$ is the number of episodes and $|\\mathcal{M}|$ is the size of the CMDP class. The algorithm is based on a combination of Inverse Gap Weighing (IGW), policy cover, and trusted occupancies and transitions. The latter is where most of the novelty lies. The regret's dependence on H,S,A seems suboptimal but, importantly, is polynomial in these parameters. Additionally, the regret is obtained with a small number of oracle calls ($\\log T$ or $\\log \\log T$ depending on assumptions). Finally, the authors present an application of their algorithm to reward free learning of CMDPs where they show that using $\\epsilon^{-2}$ samples they can obtain a model such that given a reward function, any policy can be estimated to $\\epsilon$ accuracy.\n\nSTRENGTHS:\n1. The paper resolves an important open problem (learning CMDPs efficiently with offline oracles and rate optimal regret)\n2. The general idea of trusted occupancies is interesting and may find use elsewhere.\n3. The paper has non-trivial technical novelty.\n\nWEAKNESSES:\n1. The offline density oracle is a bit vague. Can it be implemented using ERM on the log loss as in previous works? If not, how is it different than previous works such as Foster et al, Levy et al?\n2. The paper is very technical. While not a weakness itself, this complicates the presentation, especially given the limited space of the conference format. It seems that the authors made significant effort to explain the various moving parts in their work. Nonetheless, I have a few comments:\n* Why do you need to separate the data of each horizon for the model estimation?\n* Can you give some intuition regarding the definition of the trusted transitions?\n* Related to this, can you explain how the policy cover is chosen?\n* While the application to reward free CMDP is nice, I think it could be mentioned in a paragraph with the details deferred to the appendix.\n* In general, the explanations are rather technical and made it hard to build some intuition as to why things work. Even after reading a significant portion of the appendix I only have some technical understanding of how the result is obtained. I'm not sure how to address this, but it makes the paper less accessible. Perhaps showing where standard analysis would fail without components such as the trusted occupancies could be helpful.\n\nQUESTIONS:\nSee above.\n\nLIMITATIONS:\nN/a", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper studies a statistical and computational reduction from the general Contextual Markov Decision Process (CMDP) problem to offline density estimation. They propose an efficient algorithm called LOLIPOP which achieves a near-optimal regret with minimal oracle calls $\mathcal O(H log log T$) if $T$ is known in advance, leveraging a layerwise exploration-exploitation tradeoff.  In addition, their algorithm is applicable to pure exploration tasks in reward-free reinforcement learning.

STRENGTHS:
Their proposed LOLIPOP achieves near-optimal regret while minimizing the number of oracle calls. It requires only $O(H log T)$ calls to an offline density estimation oracle, which can be further reduced to $O(H log log T)$ if the total number of rounds  $T$ is known in advance. This is the first algorithm that achieves both statistical and computational efficiencies for general (stochastic) CMDPs.

WEAKNESSES:
- Some minor typos, e.g. line 466

- This paper is purely theoretical. It would be interesting if the authors provide some experimental results to demonstrate the performance of the proposed algorithm.

QUESTIONS:
- In Eq (3), to construct the trusted transitions, how do you compute $\widehat{\operatorname{reg}}_{m-1}(\pi, c)$? From my understanding, you need to compute the optimal value function $\widehat{V}_{m-1}^1(c)$, how do you implement it?

- Is there any lower bound on the number of calls for this problem?

- Compared to contextual bandits, can the authors explain more about the technical challenges of CMDPs?

LIMITATIONS:
Yes
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-13