Offline Oracle-Efficient Learning for Contextual MDPs via Layerwise Exploration-Exploitation Tradeoff
Authors: Jian Qian
Haichen Hu
David Simchi-Levi
Venue: NeurIPS 2024 poster (2024)
Keywords: reinforcement learning
contextual MDP
offline density estimation
computational efficiency
Primary area: reinforcement_learning
TL;DR: This paper presents a statistical and computational reduction from CMDPs to offline density estimation with little overhead.

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-Exploitatin Tradeoff
Jian Qian MIT EECS
Haichen Hu MIT LIDS huhc@mit edu
David Simchi-Levi MIT IDSS dslevi eeteviimit....m edu
j jiangianmiit...
Abstract
Motivated by the recent discovery of statistical and computational reduction from contextual bandits to offlin regression [36], we address the general (stochastic) Contextual Markov Decision Process (CMDP) problem with horizon (as known as CMDP with H layers). In this paper, we introduce a reduction from CMDPs 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 is known in advance. Our results mark the first efficient and near-optimal reduction from CMDPs offline density estimation without imposing any structural assumptions on the model class. A notable feature of our algorithm is the design of layerwise exploration-explooitaion tradeoff tailored to address the layerwise structure of CMDPs. Additionally, algorithm is versatile and applicable to pure exploration tasks reward-Fie free reinforcement learning.
1 Introduction
Markov Decision Processes (MDPs) model the long-term interaction between learner and the environment and are used in diverse areas such inventory management, recommendation systems, advertising, and healthcare [35, 37]. The Contextual MDP (CMDP) extends MDPs by incorporating external factors, known as contexts, such gender, age, location, or device customer interactions, or lab data and medical history in healthcare [19, 33). In an H H-layer CMDP, the learner receives an instantaneous reward at each step over 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), 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-Lee and Xu [36] demonstrates an offline-oracle- based algorithm FALCON that achieves optimal regret for general (stochastic) contextual bandits with access 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 7 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 ceaiie--rrall--ffictii CMDPs with only O(H log log T) number of oracle calls?
algorithu that achieves the optimal regret for general (stochastic)
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Table I: Algorithms' performance with general finite model class M and i..d. contexts. The optimal rate here refers to O(poly (H, S, Alg MM]). All algorithms assume realizability, so it is omitted from the table. The reachability assumption and the varying representation assumption are very stringent for tabular CMDP, for details we refer to Appendix B.
Algorithm
Regret rate
Computational complexity
Assumption
E2D [16]
Optimal
((T) online oracle calls
No No
OMG-CMDP! [26]
Optimal
( online oracle calls
RM-UCDD [24]
Suboptimal
( offline oracle calls
Reachability Varying Rep. No
CMDP-VR [11]
Optimal
( offline oracle calls
LOLIPOP (this work)
Optimal
Oloog 7) offline oracle calls
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 effffee-racleefffceenttt (as opposed to cettee-rrac---ffcieee since it requires access to an online density estimation algorithm. Foster et al. [18] provides further reduction from online density estimation to offline density estimation, with the caveat that the reduction itself is inefficient. A similar online-oracle-effccint algorithm is developed by Levy et al. [26]. separate thread of optimism-
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 ((T) number of oracle calls to the online/offline oracle respectively.
In this work, we present 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 Ö(poly(H, S, A) / [.l). This regret guarantee is minimax optimal up to poly (H, S, AA) factor [24]. The LOLIPOP algorithm assumes access to 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 Simchi- Levi and Xu [36] to adapt to the multi-layer structure of a CMDP. More specifically, the FALCON al- gorithm is divided into O(log log 7') epochs, each corresponding to an oracle call, a fixed randomized policy. However, 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 oracle call, a new randomized policy for layerwise exploration-exploitttonn tradeoff. tradeoff is done through Inverse Gap Weighting
(2) In each segment, the aeplorationeexpoiite (IGW) estimated regret for set of explorative policies. The idea of running IGW on such policy cover is proposed by Foster et al. [16]. However, their policy cover is designed for Hlyleer exploration-exploratii 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 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 near- optimal sample complexity of O (H" St A log([M)/8)/e3 with only O(H) number of oracle calls. Both the sample complexity bound and computational efficiency result for wweerd-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

---

2.1 Problem Setup
Contextual Makovian Decision Process (CMDP) is defined by the tuple (C,M - {M(c) s'),, where C is the contextual space, S is the state space, i the action space and
E S is a fixed starting state independent of the context. We focus on tabular CMDPs which assumes
S = |S], A= A|| < cc. For any context E C, (() { (c), Ri, [=) consists of H-layers
of probability transition kernel {P (c) and reward distributions (R)r(e)) [c) where P (c)
and Ri (c) are specified by Ph (- a; c) E (SS) and R (s. a; c) E A [0, 11) for all h E [H] and
E S x A. For simplicity, we also denote M = ( R hh- where Ph = {PP (c)) and R [R (c) JoEC. Let IRRS denote the set of all randomized, non-stationary policies, where any == ( ") E RRNS has x" S + (AA) for any h E [H]. We use T to denote the total number of rounds. H to denote the horizon (the total number of layers). Let M, = {PA , R he be the underlying true CMDP the learner interact with. The interactive protocal proceeds in T rounds, where for each round f, the t-th trajectory is generated as:
A context c is draw i..d. from an unknown distribution D and s} = s².
The learner chooses the policy based on the context.
For h = I..... H:
The action is drawn from the randomized policy at ~ = (s²))
The reward and the next state is drawn respectively from the reward distribution and the transition kernel, i.c., rf ~ R (s}, at ae) and A+1 ~ (- st offa))))).
Without lose of generality, throughout the paper, we assume that the total reward 0 S > م < 1 almost surely. For any model M, context e and policy æ, we use M M(T, c) to denote the distribution of and
the trajectory C12 T1, st . ai. ri. ................................ S at' .r given M, M, C1 the probability and the expectation under M(T, c) to be (.)),.((.)) and [. respectively. Given any policy и,, state s and action a, we define the action value function ( a: TT e) layer h and the value function V/th n, c) at layer h under context e and policy as
rT. Also denote
(x, = E/ (ri I of, at = 8,a) a] and V. (s; n,, c) = maxac ( a; n,, c).
We denote the optimal policy under context e as T... and abbreviate its value function as V (::c). For h 1, we further simply the notation by denoting (c) V.2 (s) and V, c) (s c).
The regret of policy π under context c and the total regret' of the learner are defined as
reg(n, c) (c) VT,(TT)) c) and Reg(T) := Σ E, (reg(m ci)l,
where [-g[-] is the conditional expectation given the interaction up to round t.
Assumption 2.1 (Realizabilitt). The learner is given a model class M where the true underlying model M, lies, that is, M, E M.
2.2 Offline Densitiy Estimation Oracles
For any model class M, general offline density estimation oracle associated with M, denoted by OfiDEM, 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 u2 by
( 1
DP(P,Q)
dP
dv,
dv
Using squared Hellinger distance for reinforcement learning is popularized by Foster et al. [16], and we adopt such 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 a distribu- tion on the set of policies IRRNS, that is, for any e E C, p(c) E A[IRNS). Given training trajectories (c,, Te, s a. ,!....,, Sll a !!! i.i.d. drawn according to G D. T. p(c) and
'The regret we defined here is conventionally known as the pseudo-regret in the literature. The conventional regret defined as Ex reg(m,, ot) can be bounded by the pseudo-regret up to an additional ((TTlog(1/)) term with a standard concentration argument, which we omit here for simplicity. "The value is independent of the choice of v.
3

---

r!. .
o .T. be the trajectory sampled according to M, ci). The offtine density estimation oracle OfDEM returns a predictor M. For any å E (0, 1/2), with probability at least 1 ò, we have
EM,((n).
The Maximum Likelihood Estimation estimator MLEM is example of offline density estima- tion oracle that achieves (m))(n) Σ log([M)/o)) (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].
Main Results and Algorithm
In this section, we present our main results and introduce the algorithm of LOLIPOP (Algorithm 1). First, we give overview of the algorithm. Then, discuss the theoretical guarantees obtained by this algorithm. Finally, we introduce the algorithm's 's different components with corresponding guarantees. All proofs are deferred Appendix D.
3.1 Main Results
The algorithm proceeds with epochs. The total number of T rounds divided into epochs. For an epoch schedule 0 To < Ti <... < TN = T//H to be specified later, the m-th epoch will last (Tm Tm-1) rounds. Furthermore, each epoch is evenly divided into H segments, each consisting of Tm Tm-1 rounds. During the h-th segment in m-th epoch, a kernel pt : C AIII) will be specified to determine the policy. More specifically, upon receiving the context C, a policy :, will be sampled from phn (ct) and executed. After collecting the trajectories fa,, mT) U {s/ a/, rt je(m in the h-th segment of the m-th epoch for Tm-1 H + (Tm mm 1) +)+1st sts Tm-1ff + (Tm mm- the offline density estimation oracle OffDEM is called with these trajectories as input. Denote the output M we will only be interested in the h-th layer of this output, which we denote by {Ph. m. Rh }. Then the collections of estimators Mn, = [Ph RA 1 will be used for the next epoch.
Overview of Algorithm 1.
T==1
Throughout this paper, we will adopt the following convention for the free variables m, म, e, h, 8, a. They will be used to denote an epoch index in [N], policy in URNS, context in C, layer index in [H), 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. T is known, then by choosing the epoch schedule Tm = 2(T/H) for m 2 1 and the offline density estimation oracle OffDEM MLEM, the outputs { ( of Algorithm I satisfies that with probability at least 6,
Reg(T) S A log(IM) log log T/8) log log T
with only O(H log To) number of oracle calls to the MLEM oracle for 8 (0,
Theorem 3.2. T is not known, then by choosing the epoch schedule Tm 2 for 2 1 and the offline density estimation oracle OfiDE = MLEM, the outputs (Te)etTT) of Algorithm satisfies that with probability at least 8,
Reg(T) S VHTSS log[[M]
with O(H log T) number of oracle calls to the MLE, oracle for 5 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 (HC/HATTM log [M] Mol/og A) proven by Levy and Mansour [24] up a poly(H) (H, H, S, A) factor.
Computational efficiency. Consider the epoch schedule Tm = 2m for m E N as discussed in Theorem 3.2. For any unknown T. our algorithm operates over (log T') epochs, making one oracle call per epoch. Thus, the computational complexity is O(log 7) oracle calls over 7 rounds, with an additional per-round cost of ((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 Layerwise pOLicy cover Inverse gaP weighting with trusted OccuPancy measures (LOLIPOP)
Require: epoch schedule 0 T-1
< TI < ... < TN = T/H, confidence parameter 8 (0, 1/2), model class M, offline oracle OfIDEM.
1: Initialize: Mo = {P) Roj Ro]h-1 where Ph is any transtion kernel and Ro is constantly 1, ... do
2: for epoch
3: Let & - 5/2n/ (mm-1 Tm-2), 2), 7m v / s /2 and 7m 7tm/720e H S A9AAa.
4: for segment h = 1, .............................. H do
5: for round
mm-1 HH- + (Tm 7m-1) 1) + 1.... 5m-1H (Tm Tm- 11) do
6:
Observe context
from the environment.
7:
for a E S x A do Compute
8:
2 (s,a;
h.s.a arg max T m,C.
SA + 7/m reg re ca)
where di, is the trusted occupancy measure defined as in Definition 3.1. Let the policy cover II,, (ca) = m } U [= 1s,0eSx4
9:
10:
Define ph, (c+) to be the Inverse Gap Weighting distribution on the policy cover I
Pm (ct; TT
VT E 1I (ct);
AA + 7m rreg-- 1 c4)
where AA is the constant that normalize the distribution.
11:
Sample and execute T. (ct) and observe the trajectory ce Tu, (5) of Jem Run OMDDEM with the input trajectories (ct, Te, (8) (e)r)))) (e(H)))) and obtain the 1-th layer estimator 5sh and Ph.
12:
algorithm [16], for instance, requires (T)) calls to an online density estimation oracle, involving significantly more calls to more complex oracle for a general model class M. On the other hand, the Version Space Averaging + E2D algorithm [18] requires O(T)) calls offline density estimation oracle and incurs a computational cost scaling with ([M) per round. Compared to our algorithm, this results in far more oracle calls and considerably higher computational costs per round.
the total number rounds T is known to the learner, we can further reduce the computational cost of as in Theorem 3.1, sim-
LOLIPOP. For any T N, consider the epoch schedule T... = 2(T 2(T/H) ilar to Simehi-Levi and Xu [36]. this scenario, LOLIPOP will run in O(log log T)) epochs, making only Ollog log 7') oracle calls over T rounds while still maintaining a slightly worse regret guarantee.
Lower bound on switching cost. There is lower bound on the switching cost of the scale allog log 2') [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 Olog log T') number of oracle calls.
Technical challenge. The main technical challenge is to accurately estimate the occupancy sures 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 single policy. To avoid this exponential divergence, we apply two methods. First, we turn layerwise design. Specifically, we generate for occupancy measures at each layer 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 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 guide exploration.
5

---

Ph m-
II Pm
11 Pm
hE(H)
P1 mi:
R
p2 R
T d
pil
m
The dependence graph of the construction. The estimation Mm-1 (Pm- R (E(A)) from the previous round provides the optimal policy 7m-1 (Line 9) and the regret estimation reg.- (Line 8) for the construction of Ie; ptr The estimation Ph. Rdh, is generated by calling the oracle OMDEM on the trajectories collected with policy kemel p/h (Line 12). The trusted transitions and trusted occupancy measures Tm d are computed from d Ph. (Eqs. (2) and (3)). The policy cover I1,s, the union of m, and the policies (TT1):a) ( J calculated in Line which requires Th-11 Thh The policy kernel pmn is the inverse gap weighting on I (Line 10).
Figure 1:
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 I). We first introduce how the estimators {Ph N from the previous epoch are used in the new epoch. Then we proceed to introduce how to construct ph (c:) given 11th (c+) Next, we introduce how are ma P obtained given ph: Subsequently, we present the definition of the set of trusted transitions T and trusted occupancy measure d Finally, present how the policy cover IT (cx) is constructed.
During epoch m, we will be using the estimators 1Ph-1. - from the previous epoch for regret estimation. More specifically, for π,, c, h, 8, we define the value functions with respect to the model (Ph-1(c), (c) as V (8; c). The optimal value function is denoted by V1 1 = max, 1 (3: (s; c). For h = 1, we further simply the notation by denoting (c) Vn-1 (s) and VI (IT c) (s'; c). Also denote the optimal policy under context by 7m- = arg max, PI
1H
m (T, c). Thus, the regret is estimated to be
reg ( (c) - P (T, e).
At round t, let m(t) and h(t) be the epoch in which the segment round lies. We note that during each epoch m and segment h, all of the notions IT,, (c)), Phr (ct), R (c),, (c)),, d (e)) h.s.a will not depend on the specific time step i, but only the context a. Thus, we
ph (ct), and mhee, will use 11tt (c) denote the policy cover if c; when m(t), h(t) = m, h. Similar conventions regarding the context e apply to the notations Ph. Ph. TA
m. dh phh Tm. Under any context c, the policy cover II will include mm and has no more than SA + policies. These two properties together guarantee that the Inverse Gap Weighting [14] randomized policy p/: (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 ph (c) is well defined, i.e., there exist No. E [0, 2S AS such that Exemhh (o) Pm,..()) = 1. Furthermore, we have the estimated regret is bounded by Exepph (c) ( S VH A..
The choice of A/... here is for compactness of presentation. It can be chosen to be 25A for suboptimal arms and collect the probability remained to the optimal arm [36], which is computationally efficient. _((t),saa for any t, s, a can be computed in poly(H H, S, A, log T)) time
The computation for the policy *m(())c; by formulating it as a linear fractional programming problem. We defer the details to Appendix G.
Since / maps C to randomized policies, it thus policy kernel. This means the trajectories gen- erated in epoch m and segment h follow an i.i.d. distribution as described in the definition of Defini- tion 2. I. By applying the guarantee from Definition 2.1, we have the following guarantec on m. R

---

Lemma 3.2. For any m, h, and c ~ D ~ (c), we have with probability at least 2N that
8
Bar MM [Di (Ph (s af:c) C) (s af:)) + D2 RRh (St, at: c), (s) af: af:e) e) - E
If the offline density estimation oracle is chosen to be the Maximum Likelihood Estimation oracle MLEM, we will obtain &m S /
Tm-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, a, c, h, 8, a, we define iteratively the trusted occupancy measures dh (s; NF. c), dh, (s, a; (s,;; and the set of trusted transitions Thi (c) at layer h as the following:
(s5T,))) 1(s= 2n (s; = -M I' (,) (sss'sss',
(s, 1(s =s') (s)
(2)
as c) ==
For any m, h. the set of trusted transitions are defined as the set of transtions
Th (s, 18,0; aT
(c)
(s,a,5)) max SA+ 7mm reem-1(TT,)
ph ( 2
(3)
where Sm =
"ho SeH(H+1y7 Notice that to define d (8; ()) and d (s. 0; π. c). we only need (c), PL, (c)3jE(n-11 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.,
(s; 1(s==s), (, 1(s 1(8= s''=''''ssss
(s;T,)) : ( (a) (s', (sls',
(s,0,, :d (s,a).
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 7) time by formulating it as a linear fractional programming problem. We defer the details Appendix G.
Define the estimated occupancy measures di (s; nT, R. ( x)] and d (s, a; x) ==
B.M..t,ee1 (sft at 8, 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, n, h 8, a, under any context c, we have
Pi (s, v.C) Tr. (8,0; c) 32e//H55 &m + (m;
The next guarantee is the key to our amalysis and is the most non-trivial guarantee our construction. The following lemma states that if, for a context the Hellinger divergence between P and ph at layer is small for all h E [H]. then the trusted occupancy measure is upper bounded by scaling of the observable occupancy measure.
Lemma 3.4. For any m and e, assume for all hs [H].
Euneh ( M.i.,e
D2 (ph (M
H/mm-
Then for the same m, e and all n, h 8,0, we have
(s, S 1 (s,a;
7

---

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 1 and pmn Concretely, we upper bound the trusted occupancy measure d (,,;),))) of any policy и by the trusted occupancy measure induced by policy kernel pin.
Lemma 3.5. For any m, C, h, 8,a, have Z (s, (s,a; C. Du(P" (s, Pm (s,
- 7m
Pm (c, mi h.s.a) )d (s, a; ame ,c) (s 0; Ph (h + En
(4)
where En
20²
Em H
+ S r0 1 (T, d (s,0; ,and p (c, aia is
the probability of pl)(c) Thea. The guarantee Eq. (4) also holds replacing (s(P. (3, ph (h(ss,
()) with Du (s a; c), Rh (s,a;a))) on both sides of the inequality.
4 Regret Analysis
In this section, we prodive 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 to present the following epoch-wise guarantee.
Lemma 4.1. For any any policies {TeeeEe and ā E (0, 1/2), with probability at least 1 - 8/M,
EOND | Vinl (T..) - V. (Fez | S 1 EEND regm ( c)] + 77ev 77 A3 E
20
Proof sketch of Lemma 4.1. For simplicity, in this proof sketch, we assume the reward distribution is known". 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,
m (Te, (Te, c) (MTe c) Σ h...a
(s, ne, Du (P (s, as; (s,aa)))
Then we can exchange the estimated occupancy measure d (s, a; ae, c) by the trusted occupancy measure through Lemma 3.3, that is,
Zi (s, a; Te: DH P (s, a: c), P), < Th (s, a; nec DH (s, (*(,
3232 H A Em (
Then by coverage guarantee Lemma 3.5, for any h, s, a, we can bound
Zh (8, a; Te, c) Du
phh (s, c), P. (s,
< 7m
h.s. Pm (c, mm ) ( (s, (8,0; h.s,a ,)) (HPP (P. (s, c), Ph a; e)) + En
If the assumption in Lemma 3.4 satisfied, then by Lemma 3.4 and the definition of ph have
Σ 7'
p (c, whhisa )dt a; hhs c) (P (s, a),, (s,
eH
h.s,a
7'
Phi (c, o Th )a (8, oz m. c) D2 (PP c), Phh (s,aac))
haa
H
b,s.a
7 H
Σ Enmph (c) M. a. [Di (fih)) aft)) of
<
'The uncertainty in reward distribution is not the main hardness in this problem. Specifically, our proof extends to offline oracles with squared loss guarantees between the mean rewards for divergence between the reward distributions [8].

---

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
BooD
V (Te, c) Vi (Tee e) S E r0gm- e) /40 + 39ev H Em-
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 found in Appendix E. Combining Lemmas 4.1 and I.I, we arrive at the following general regret guarantee.
Theorem 4.1. The outputs (Te)eeTT) of Algorithm I satisfies with probability least 1 ō that
RTg(T)) $ Σ mm (Tm Tm-1) 54 Em
for 8 E 1/)).
5 Extension: Reward-free Reinforcement Learning for CMDPs
In this section, we introduce the application of Algorithm I in the task of reward 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 to output a CMDP prediction M whose value functions are close to the underlying true CMDP M, for any reward distributions. To accommodate such change, for any model M = {P R& ]H[HH] and reward distribution R = ={ {RR Jee(t)) we define M (-; R) = {Phr Rhh M = (P(())). (c)) (C)))))]))) ceC 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. = (-; () where 0 is used to denote the reward distribution that is constantly 0.
For any model M reward distribution R. context and policy distribution of the trajectory C1, TT,
use M(T, c; R) to denote the $ $1 a² ri....,,s off of .ri given M. = M (-; R), C1 = c, and Ti = n. Also denote the probability and the expectation under M(x c; R) to be R.a....())) and EM,,,,,R(() respectively. Given reward distribution R, any policy म, state s and action a, we define the action value function Q*: (8, a; n. c, R) at layer and the value function V7 n. C, R) at layer under context and policy क
H
(s, R)
Σ
BMM. M. sst,, at = s, a] and (s,T,,, = max (s, R). aEA
We denote the optimal policy with reward distribution R under context c as T,,,, and abbreviate its value function as (:;e,
e, R). For ht = 1, we denote V22 R) = (s R) and (T, R) V I., c, R). We also denote V 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 E -M have 0 reward.
For given E, 8 > 0 and model class M, the goal of the learner is to output model M at the end of the interaction such that for any reward distribution R and set of policies (Te)aecs the model satisfies
EADD V7 (TT, (T,,, c, R) - ViJ (Te, c. R) S
(5)
with probability at least 1 - 8. An algorithm that achieves this objective is called (e, de-rrss the model class M. Then we have the following guarantee from Algorithm 1.
9

---

Theorem 5.1. If we choose T/(2H) and T2 = H, the outputs M2 of Algorithm satisfies the reward-free objective Eq. (5) with probability at least 6, with T at most bounded by
T O St A 1og([M
for 8 E 1/2). Moreover, the algorithm requires O(H) number of oracle calls to the MLEM oracle.
The proof follows similar argument of Lemma 4.1. In addition, have matching lower bound up to poly(H. H, S, A) factor adapted from the non-contextual lower bound from Jin et al. [22].
Theorem 5.2. Fix S 1, 1/2, A2 2 Suppose S 2 Log A for large enough universal constant I. and K 2 O large enough. Then, there exists a CMDP class M with [S] = S. |AA A, [M] K, and horizon H and distribution μ on M such that any algorithm Alg that (E (e/24, 6)- learns the class M satisfies E E M. (T) Z log |M]ee,, where is the number of trajectories the expectation under required by the algorithm Alg to achieve (e/2A, 6) accuracy and M, the interaction between the algorithm Alg and model M.
Discussion
Extension to low rank CMDPs. Low rank MDPs represent significant extension tabular MDPs, as explored in various studies [6, 28, 21, 4]. Linear MDPs are typically the first step beyond tabular MDPs. Extending our approach linear CMDPs would be 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 does not utilize the low-rank structure.
Extension 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, opposed 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 explore whether reduction from CMDPs to offline regression could be feasible.
Acknowledgements
We thank Alexander Rakhlin, Dylan Foster, Yanjun Han, Zak Mhammedi, and anonymous reviewers for useful discussions. JQ acknowledges support from ARO through award W91INF-21-1-0328 and from the Simons Foundation and NSF through award 203-2031883.
References
[1] Alekh Agarwal. Selective sampling algorithms for cost-sensitive multiclass prediction. In International Conference on Machine Learning, pages 1220-1228. PMLR, 2013.
[2] Alekh Agarwal. Miroslav Dudík, Satyen Kale, John Langford, and Robert E Schapire. Contex- tual bandit learning with predictable rewards. In Artificial Intelligence and Statistics, 2012.
[3] 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.
[4] 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.
[5] Philip Amortila, Dylan J. Foster, and Akshay Krishnamurthy. Scalable online exploration via coverability, 2024.
[6] Steven Bradtke and Andrew G Barto. Linear least-squares algorithms for temporal difference learning, Machine learning, 22(1):33-57, 1996.
10

---

[7] Abraham Chames and William W Cooper. Programming with linear fractional functionals. Naval Research logistics quarterly, 188 1962.
[8] Fan Chen, Song Mei, and Yu Bai. Unified algorithms for rl with esttimatinnn coefficients: No-regret, pac, and reward-free learning. arXiv preprint arXiv: /745, 2022.
[9] Jinglin Chen, Aditya Modi, Akshay Krishnamurthy, Nan Jiang, and Alekh Agarwal. On the statistical efficiency reward-free exploration in non-linear rl. Advances in Neural Information Processing Systems, 11.20960-209777 2022.
[10] 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.
[11] Junze Deng, Yuan Cheng, Shaofeng Zou, and Yingbin Liang. Sample complexity characteriza- tion for linear contextual mdps. arXiv preprint arXiv:2402.0P 02700, 2024.
[12] 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.
[13] 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 Artificial Intelligence, pages 169-178. AUAI Press, 2011.
[14] Dylan J Foster and Alexander Rakhlin. Beyond UCB:
Optimal and efficient contextual bandits with regression oracles. International Conference on Machine Learning (ICML), 2020.
[15] Dylan Foster, Alekh Agarwal, Miroslav Dudík, Haipeng Luo, and Robert E. Schapire. Practical contextual bandits with regression oracles. International Conference on Machine Learning, 2018.
[16] Dylan Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2 13487, 2021.
[17] 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.
[18] Dylan J Foster, Yanjun Han, Jian Qian, and Alexander Rakhlin. Online estimation via offline estimation: An cotommaiionthhorettee framework. arXiv preprint arXiv:2404. 2024.
[19] Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes, 2015.
[20] 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.
[21] 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.
[22] 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.
[23] Yin Tat Lee and Aaron Sidford. Solving linear programs with sqrt(rank) linear system solves. arXiv preprinl urXiv: 2019.
[24] Orin Levy and Yishay Mansour. Optimism in face of context: Regret guarantees for stochastic contextual mdp, 2023.
[25] Orin Levy, Asaf Cassel, Alon Cohen, and Yishay Mansour. Fluder-based regret for stochastic contextual mdps. arXiv preprint arXiv:2211.14912 2022.

---

[26] 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.
[27] Gen Li, Yuling Yan, Yuxin Chen, and Jianqing Fan. Optimal reward-agnostic exploration in reinforcement learning. 2023.
(28) 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.
(29) 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.
[30] Zakaria Mhammedi, Adam Block, Dylan J. Foster, and Alexander Rakhlin. Efficient model-free exploration in low-rank mdps, 2024.
[31] 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.
[32] Aditya Modi and Ambuj Tewari. Contextual markov decision processes using generalized linear models. arXiv preprinl arXiv: (903:069900066666 2019.
[33] 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.
[34] Yury Polyanskiy and Yihong Wu. Lecture notes on information theory. Lecture Notes for ECE563 (UIUC) and, 6(2012-201)) 2014.
[35] Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
[36] David Simehi-Levi and Yunzong Xu. Bypassing the monster: faster and simpler optimal algorithm for contextual bandits under realizability. Mathematics of Operations Research, 2021.
[37] Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[38] Andrew Wagenmaker, Max Simchowitz, and Kevin Jamieson. Beyond no regret: Instance- dependent pac reinforcement learning, 2022.
[39] Ruosong Wang, Simon S Du, Lin Yang, and Russ Salakhutdinov. On reward-free reinforce- ment learning with linear function approximation. Advances in neural information processing systems, 33:17816-17822 2020.
[40] Yunbei Xu and Assaf Zeevi. Upper counterfactual confidence bounds: new optimism principle for contextual bandits. arXiv preprint arXiv:2007.07876 2020.
[41] Yunbei Xu and Assaf Zeevi. Upper counterfactual confidence bounds: new optimism principle for contextual bandits, 2024.
[42] Xuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang. Alekh Agarwal, and Wen Sun. Elficient reinforcement learning in block mdps: A model-free representation learning approach. In International Conference on Machine Learning, pages 26517-26547. PMLR, 2022.
[43] Zihan Zhang, Simon S Du, and Xiangyang Ji. Nearly minimax optimal reward-free reinforce- ment learning. arXiv preprint arXiv:
2020.
[44] 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.
12

---

A Notation
For any integer n, we use [n]] to denote the set (I....,m) n). For any set X, we use ( to denote the set of all distributions on the set X. We define O()), 2()), (-)), Ó()) 2()). θ()) following standard non-asymptotic big-oh notation. We use the binary relation x S y to indicate that z 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 con- textual bandits with access to an online regression oracle. This is extended to the CMDPs by Foster al. [16] with the E2D algorithm. However, the algorithm requires ((T) called to an online density estimation oracle. Compared to our algorithm, it necessitates significantly more calls to an oracle that harder to implement for a general model class M. After the FALCON algorithm [36] es- tablishes 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 Pmin to all states under any policy and the regret guarantees scale with O(poly H, S, A) (1/Pmin) / log M[). This assumption precludes model classes with small reachability, which frequently happens in practice [4]. The CMDP-VF algorithm proposed by Deng et al. [11] assumes a varying representation assumption on the model class instead. The assumption asserts that any model M (Ph, R* NE(H] E . satisfies for any e, h, s, a, s', Ph (s'Is, a; c) (o" a; a; p² (s')) for the known feature vector (s a; E Rd and an unknown vector (s") Ted 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 8, a in RSA, i.e., d (s, az c) Cz.a S.S.A Then, the requirement of p²(s' not depending on the context forces ph 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 (poly (. d) TTVTloog [M). Another significant disadvantage compared with our algorithm is that the RM-UCDD and CMDP-VR algorithm requires ((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, 5J. However, to the best of our knowledge, we are the first to study the reward- frre reinforcement learning setting for stochastic CMDPs. We provide a near-optimal sample complexity upper bound and matching lower bound up poly (H, S, A) factor with only O(H) number of oracle calls. Nevertheless, the upper bound is obtained by adjusting the exploration-explooittioo. highlighting the fiexibility 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 finite model class and the MLE estimator M be defined by
M arg max MEM i=1
p Ma (1s) r heem))
For any 8 e (0, 1/2), we have with probability least 1 1- 6,
log(IM1/8)
E BewDe
D M (( c), M. (T, x
n
13

---

Proof of Example C.1. For any 8 E (0, 1/2), by Lemma C.2 of Foster et al. [18), we have with probability least 1 8/2,
Σ Di M(TT, M. S log(I
Then by Lemma A.3 of Foster et al. [16], we have with probability at least 1 6/2,
Ba-Derolo) pi (M, o), M,, ~ Σ Di
+ log(1
Then by union bound, we obtain the desired result.
C.2 Information Theory
Lemma C.1 (Lemma B.4 of Foster et al. [17]). Ler P. and Q be two distributions space Let X - be a function. Then we have:
1- - -el S (2
C.3 Reinforcement Learning
Lemma C.2 (Lemma F.3 of [16]). Let M = {PM, Rhy heet and M {Ph Nr & heem be two
CMDPs. For any policy æ and context we have
(Tr c) v (T)
H
E M. Te (PM A+1 1s,,
hml
H
E m,
E Bnd R (1
h=1
H
( (s, M, M. [ at = =s (DH (s, (;),
+ Du (a( a),, (f, (s))))
1=1 8,0
D Proofs from Section 3
In this section, we present the proof' for Lemmas 3.1 to 3.5.
Proof of Lemma 3.1. We fix an arbitrary context throughout the proof. Let u(A) ETEnA. 1/(A + Tm (egm-1(T)). Since for any h E [H]. Am-1. E I then u(A) 2 1/A. On the other hand, u(A) < (SA 1)/X. Moreover, u(A) is clearly monotonically decreasing with u(0) oo and uSSA + 1) & 1. Thus there exists Ah. E (0, (A,SA+) 1] such that u(>an.)) = as we desired.
1/(X + 7m (e8m-1(fm--id
Now we have the regret is bounded by
Exouth (c) [regm-1f [regm-I(TT()) =
Σ
reg (T;C)
2SA
A
+ Um re8m-1 (T:C)
Tm
TEIIA
WEIT (c)
Finally, recall 7m = Tm 7m/ (mm A2) and Ym
4o Em
plug this bound then we have
E B (c) freg (T; < 1440e3.. 11 Em.
Proof Lemma 3.2. By definition of Definition 2.1.
14

---

Proof of Lemma 3.3. For any m, c, h, s, a, by the definition of d (s, a: c),, we the difference
between d (s, n. c) and d (s,a; c) are the parts of occupancy measures that do not go through the trusted transitions, i.e.,
a (s,a; TT,)
din (s)0;
h-1
Σ
Σ
Lnlss, (als),a;T, (s (s
nT, (s,)),
(e) a.a a (T)
where P1 (s]s² +1 ग,, c) is the estimated transition probability from sht at step j + 1 to s at step h under policy क and context Then since (s), a² s & (c), have
h-1
Σ
(s a n c) P/ (ا ( a²; c) Paitt (s)s cs)"
m+1
j=1
(ef al
1 (m))
Σ
L a² nr, c) ( ( ,
(0) ./- (
SA + 7m mm
Σ
Sm
1-1 (ef (0, (c)
hS2 A(SA Teg (T:
Sm
Recall the choice of Sm
H 7m
7m
720e A
and 7m
V
11
Em
we have
(s,00;,, - ( &m
Ieg - (TT; c).
90HSA
Proof of Lemma 3.4. We prove iteratively between the objective
(s s (1 11) 2(h-1)
(s
and the following claim: For any h and any ss) E (c), we have
1
(s S
1 +
p (s'1s,,
H
First, we show that for any h E [H].
1 2(h-1)
Va (s < + H
v's,a,s)) E (c))), Ph( S (1+ H
1 1)² (s(ls,acc)) (6)
For this, note by Lemma C.1 that for any h, h,
Ph (5(1),));)))
+ 1 (2 PM (C) Ber (s. {P((ss)) + (,lsssss;;) PP
(8)
15

---

where the second inequality is by data-procese processinn inequality [34]. Then by AM-GM, we have
2 (Ph( (881),)885) + PM a;) Di
Pih (8, 0z Ph PM (b,
(Ph (5' 8, aze) (s + Plug the above back into Eq. (8) and reorganize, we obtain
+ Dí phh P,
S 4H
P & [ (H) Ph (s'ls, + (H 11 (PM(,,a))) (M)(((((ssrrr))))))....))))
Then multiply both sides by din (s, vT c), we have
Pt (11 ((41))6111 < (H (s, (P)
Im (s, (s, (o' 05c) (1 Meanwhile, by the definition of ahes we have
Zr (8, a; m mme c)
h.a.a
h.a.a an (s, SA tm regm- (Tm, ( c) Then by the induction hypothesis, we have
(SA + Tmm (8,
I h.s.a The (8, m.c ,)
(SA+ + 7mm reege- (18m-1(FT,)))
SA + 7m. regm- (nthi ,c)
ed (s, 0j m,e
(SA + 1m reg (
SA + Um regm- (Tm. .c)
Thus
p/n (c) and the assumption that E (c) B.M. Di pm (s a); (* at:c)] s H/ 7ms
we have
futher by the definition of
(11 1)dm (s, (s, D) (s c))
dh (s, a; whes c)
S S + +)) 7m rég ( SA + nm regm- ( h.s. e)
Dj (hh (s, P,),a
5 + 1)((AA+ + Vm regm (TT, c) Pm (c, h.s d (s, aj mme De (s, Ph
2H + 1)(SA+ + Vm reg (TT, c) Bennit (c) M. Di Ph af) (sf.
H(H +1
(5A + 1mm regm (TT, c))
7m
(SA + 1m regm (5,,)((,e)))
(H 1)6m
where the last equality is by the definition of Sm- In all, we have shown that
Th (s, a, e) (s - )
(h (111..))))
S (H 11Sm (SA + 7m reg.-1 c)).
1
Now we prove by contradiction, if for any s') E (c), the reverse inequality is true, i.e., ph (s'|ss (1 + p) Then for any
ph (h((s')s, (1
H + The (s. 0; n. c) Ph (8'1s, a3 < 2hh (s, a; n,
ph (s' (s'1s, a;c)
1
< H 1)
(SA+ + 7m regm cT).
16

---

This contradicts the definition of Tth (c).
For the other direction of Eq. (6), we prove for all h E [H]]
Ps < (1 + H 2 (h Thh (s, & (1+ #)
a(s,a,s) E T
2(h d (s, VT,
2h
2+ (8,0; =
1
d (s, Va,8,a.
(9)
H
This direction is straightforward since
g (s,
Σ
a c)
(s', TM
2h
+
H
1
2h
d (s,
H
With the two derivations Eq. (6) and Eq. (9), along with the fact that the initial argument Eq. (6)
holds by defintion for h = 1. Thus we conclude the proof.
Proof of Lemma 3.5. For any m, s,a, by AM-GM, we have
Tn (s, a
a, n, c) (s, Ph (s)
2e HSSA + 7m (TTh (Th.s.a..))
(2,
Im (8, as Tm,c ;
(10)
You dm (s, mmm =bbsaa;))
Di (s, a; (
2c HSA 1m Fegm- 1 (omaen:c)) :))
By the definition of mhe we have futher
2SA + 7m 10gm- 1 (Thie c))
(d (s;
Ymm d (s, a; has,a ,c)
2c HSS++ mm
(11)
7m (8, C)
252HSA + 2c Hmm rog 8m (T;C)
dr (,)
Yin
Recall the choice of Tm = 7m 7m200 53 and 7m
M484 Em
we have
2 + 2c Hmm (Tzcc) 26²
Em
regm (T)
(12)
H
8 A2
Also by the definition phn (c), we have
Tm To (8, (z; o n.s,a
D2 (p) Ph (h(())
2e HSS++ + reg - (Tiie (c))
(13)
$ Ym
pi (c, hisa )dn (s, a; bis D) (s,,,)))), (P) PM(,s,)
²H
17

---

Now we plug Eqs. (11) to (13) back into Eq. (10) to obtain that
(s) (s,
Phi (s)
&m HAS2A
20
720H4 reg 1
I (s,
mm e
P'm (c, mm ya (s, a; mm c)
Ph ())
Similar
bounds
can
be
obtained
replacing
Du (P ()) (s.
with
Du as, a; (s,a,))
E
Proofs from Section 4
Proof of Lemma 4.1. For this we first apply the local simulation lemma (Lemma C.2) in rein- forcement learning to relate the divergence of the value functions to stepwise divergences as the following. Under any context
V TT. c) V. (Zc. S Σ as ne; c) (Dn
(s, (s, (P Dn
RR (s, (s, (s,
h.8.0
Then we can exchange the estimated occupancy measure dn (s, a; Te, e) by the trusted occupancy measure through Lemma 3.3, that is,
Σ 2r (s, Tes h,a,a
DH ce, (P (P)) P.
dn (s, a; Toy DH
(P P
h,s,a
+ Σ
32ev H+ Em Em + 90HSA
1
reg 10gm-1 1 (T:
h.s.a
dh (s, a;
Tey DH (sh a: Ph (s,
h,a,a
32c A3 En +
regm (Te; c).
(14)
90
Then by coverage guarantee Lemma 3.5, for any h, s, a, we can bound
Σ
(s Te, DH (s, (s,
h,5.0
7'
phn (c, mhhsaa) (s, a; whesa .c) (P (h(ss (h)
e²H
h,3,0
Em
Σ h,s.a
2e²
1AS2A 7 A2 regm - (T) (T Zh (s,0;,,))))
Σ
7m e2
pith (c,
.heaa ) (8, ar m
b.s. .c) DP (h) (h)
Em
2
reg 109m-1(TT)) (Te,
(15)
H2
720 A
18

---

Suppose the assumption in Lemma 3.4 is satisfied, then by Lemma 3.4, we have
Σ
7
h.s, Pmmle, (c, mm dm (s, as Ti h.s.a c) (s,
Pmn
e²H
h,s,a
V
"mm(s (c, "ieee)) (s, mm c) ( (8, c), Pm (s,
b.s
H
b,s,a
Ym
E Bxppph (c) B. Dit (P af:c) (" (s:
(16)
H
The Du ( afc c), P. af;)) in Eqs. (15), (16) and (26) can be replaced with DH Rh (sf ofic)) R (s) afic) as well. If the assumption in Lemma 3.4 is not satisfied, then we there exists such that
Ba (c) (o) M.ma [Di ph (s) P. PM(s > H/mm-
This implies
m (Te: c) - v. (Te, S 1
mm
-B (6) BM. ai Di (o p2
H
Thus, altogether, no matter the assumption in Lemma 3.4 is satisfied not, we have shown
Vi (Tp, c) Vi2 (T o) s 76c (6Ge/H55*Af A3 Em +
1 1
20
2mm
Σ
E (p)
M.,...
Dit
ph (st, af;af Ph (ftt
Taking expectation on c, we have
17 (Te, c) V77 (
BoND
EonD I + 76ev 1654A3. -E
20
+
2m
Σ
EeNDEarpphe
M Di af o
H
h
1
E reg
(Ter c)l
+ 76ev 43 Em 7m &m
20
1
E reg regm-1
( c)] + 77e 43 Em.
20
where the second inequality is by the offline density estimation bound from Lemma 3.2
Lemma E.1 ([36)). Let E1...., Ex be N positive values. Suppose for any m > 0 and an arbitrary policy set (To we have:
E [1 (TT, (m,, S EowD ( (rögm- 1 + Em-
(17)
20
Then for any m 0,
10 E.-Dlreg (Tey c)] 8m,
E c)]
(18)
9
BonD Feg
(Tes c)l
S
BenDlreg(Ter c)] + 8m
(19)
8
where 81 = 25, + 1o and 8m f/ +
20
m for any 2 2.
19

---

Proof of Lemma E.1. We present the proof here for completeness. By Eq. (17), we have that for all m and Te
BowD/reg( c)] EewD fregm c)]
EewD IV ( (Te,, c) - 171 m (T c)] + BewD (V (T - V (fm.
+ BoND (V (T c) (TTs c)]
< EaND reg (T c)] +
Bew re (Tes c)] 28m
(20)
20
20
Symmetrically, we have
BowD regm Tes c)] Eo-p(reg(Ter c]]
E EEND feem- (o c)] + Bewollegm (Tes c)] + 22
(21)
20
20
Now we inductively show for 1,2,... that Eq. (18) and Eq. (19) hold. For 1, since Bewp/reg(TTA c)], BewD regı (Te, c)] I for all 9,, then we have from Eq. (20) and Eq. (21)
BeED (reg( Tez c)] R [reg, (Te, c)] + 281
and
B [regi Te: c)] < RewD (reg Te, c)) + 2s1 +
10"
Hence we have shown Eq. (18) and Eq. (19) m 1 and 2s1 + 10 Now, suppose Eq. (18)
and Eq. (19) holds for all ...,..., 1. Plugging Eq. (19) for m Eq. (20), we have
1 into the right hand side of
EO (reg( Te, c)] BOWD (Ter c)]
Eew reg (T (T)
c)] + - Br- c)] + 2
20
20
<
1
9
BewD(reg((TT.) 6)] 8m 8m
1
9
EunDlreg(Tes c)] 5m-1) + 28
20
8
20
1
(
8
+
Brew Te. c)] + 8m-1
+ 25m:
20
20
where the last equality is by Be-D (reg( TT.., c)] 0.
Thus reorganizing the terms, have
10
c)]
E EenD reemm (Te. c)] + 8m:
9
where 8m = fmmm1 m. Thus we have shown Eq. (18) for m. Then plugging Eq. (19) for m into the right hand side of Eq. (21), we have
EonD (reg.. (Te, c)] - BENDD (oo.D(rrg(Tee c)]
1
Enn (regm ( c)] + Bew faeed-reemmae ( c]] + 28m
20
20
1 9
(9 BewD|rog(Ter e)) 8m-1) + 2mm-
Peww/reg(fTma.. c)] 8m-1)
<
20
8
20
10 Furthermore, by Eq. (18) for m we have Ee- (regg(ff mo c)] S BowD (reg c)] + 6m = 5m- Plug this in the aforementioned inequality, we have
7f
Ee-D[feg,, (Te, c)] -...eDrrrggaa.. - 1 fm + 8m-1) 20
20 Bewplrog((Tr c)l 6m-1) 25 m-
1 9
Reorganizing the terms, this in turn gives
Ean (reg (Te, c)] S
c)] 5m,
20 where recall 8m fmmmm11 +
18 20 This proves Eq. (19) for which completes induction.
20

---

Proof of Theorem 4.1. Let & = Then we have by Lemma 4.1, Eq. (17) holds with 6m = AV for L > 0 large enough for all m O. Combining Lemmas 3.1 and E.I, we have
(oleerr bo (c) (reg(m, c)] X BewD,r (c) (To, c)]
9m V HOSIAS.C &
Then by Lemma 3.1, we have further
Bew (6) Teg ( (T)
Hence, we have
Dow (c) (reg(T e)
9m H35443. &
In all, we can obtain the following regret bound with probability at least
RTg(T)
Σ
B FooDD D (aa (cr) [rog( cil
Bew D p (r)e(((T (Tm Tm- 1)
h.m
N
1mm-Tm-l ....).HHHS. /HSSAA3 Em.
mm
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 choosing
Tm 2(T we first note that since Ty S H/H, we have (T/H)² 2 2. Then we have
T 2(T/H) 1-2 2 1-2
= -2(T/H)1-2- 1- (T -1)
2 2(T/H) Tm-1-
This implies Tm Tm-1 2 2 Tm for m [N]. Then by Theorem 4.1, we have with probability at least 1 6,
T
Reg(T)
Σ (Tm Tm-1) A4
$ (o
5m Tm-1
+ Ti
43
Tm-
Tm-2
S / 49
T'm
+ A
Tm-1
S H 1og( VT/H
VT/H H LogMM/Tooo T/8) log log T,
where the last inequality follows from
2(T/H)
Tm
(1-21-m = 2
mm
(m 1)/2
and N O(log log T). So the number of oracle calls are O(H log log T).
21

---

Proof of Theorem 3.2. Since we are choosing 2", 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 6,
N
Reg(T)
S
(Tm- - Tm-1) 13 Cm
m=1
N
2m-1
X v lo (M 1+ Σ
V2
2/2 m=3
x v 1115549 M M//00)
/4TS4 log(I log T/8).
Proofs from Section 5
Proof of Theorem 5.1. The reward distributions R., are 0 for m 0, since Ro are set be 0 and the model class M consists of models with constantly 0 reward. The regret estimations regm-1 are all for m 2, Thus we apply the component-win guarantees (Lemmas 3.2 to 3.5) P2, Rs where R₂ = 0. Concretely, from Lemma 3.2 we have with probability least d,
D. M......
p Ph (s) P.
E2.
(22)
with the offline MLEM oracle. From Lemma 3.3, we have for any C, s,s, (s, agm, - (s, a; a; S 32e 32e//H522 H A Ez. From Lemma 3.4, have if for context and all h E [H]]
(23)
E (c) BM
[ P afc)), P. af: < H/72
Then for the same and all n, h, 8, a, we have
(s, S (1+ 11 21 d2 (s
(24)
Finally, from Lemma 3.5, we have for any n, n,
as, DH(P (s, P) (s,
< H2 (e
T2. (
hea
و (s, a; T2 mhe .c)
(25)
Eg
+ 2
(s, a)
H
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 {Neloec and any context c, by local simulation lemma (Lemma C.2)
(T R) Ps (T c. R)
S Σ
Ts a; Te, Dn P (s, Ph (8. (s.
h.e.a
Then by Eq. (23), we have
Σ h,s,a
3y (s, a; o
DH 1h (s, (s, (;)) (P)(s,ac))
Σ (e (()
(s + 32e, 4 &g (26)
h.s,a
Case I: If for a context e and all E [H].
E Eanopy (c) M.
D Ph (4, af;e,, (ff. < H/72
22

---

Then by Eqs. (24) and (25),
Σ (s, D)
(1)
e²
12
p2
P2
2.
h.s.a
as, her .c) (3))
E2
2e²
(3,0
H'S2A
12
H
P2
h
(c,
22
h.s.0
d
(s,
hee .c) (s, ((,,))))
E
2e²
(s.
H
Thus have (T R)- (Te, C R))
x VHOSTAS.E.
E (3 M.
Di Ph. (sh, P (ot
H
h=1
Case If for context e there exists j E [H] such that
E xa (c) BM Di (P) > H/2.
Then
(T,, R) (Te RD)
11
72
ampp (c) B.,... [Dt (P)(P1 P2
1 <
H
1=1
Combine Case and II, we have
(Te., R) - Vig (Te
$ &
72
Σ E pp
E
M.
Di
Then take expectation with respect e together with Eq. (22) obtain
EEWD V. (Te: c, R) - Vz (Tes e, 1| x HoSSA3.. &g + 7-82
H7 A3 (M
24
T
24
8,
where the second inequality is by Eg H log( and the last inequality holds when
T Ω 5 A log( MM
Thus the reward-free objective Eq. (5) is satisfied with probability at least 8 with
HJ A (lgMA)))(())188)
T
2
23

---

Lemma F.1 (Lemma D.2 of Jin et al. [22]). Fix E I 1. 8 S 1/2, H.A 2, and suppose that S Log A for universal constant L. Then consider the trivial context space C {cob. there exists model class M {Mj]jeJ, with |8| = S. | = A, |.1 S eSA and horizon H and distribution i on M such that any algorithm Alg that (e/ -le the class M satisfies
SA
EMAIEMAA E²
where T is the number of trajectories required by the algorithm Alg to achieve (2/12, 6) accuracy and E Algl [-] is the expectation under the interaction between the algorithu Alg and model M.
Proof of Theorem 5.2.. Let n. log (/(SA), then consider the context space C {c1, --1.- Cn and the i.d. distribution D on the context space being uniform. Denote the model class obtained from Lemma F.1 by M = {M,(co))jes
Let J ...,,.. In E J be an index. Then we consider the model class M = {M,]JET where Mc(c)) = M, (cn), that is, the model class M is on each context C E C an independent M. We first have that the size of the model class is bounded by [M | " I eSSA <. K. Then have for any algorithm Alg that (#/24, 8) bllaarsss the class M, must have (e/12, le-lrs the class M(c) = (M)(G)))eee>
for least half of the contexts E C by Markov's inequality. This, in turn, combined with Lemma F.1 gives that there exists distribution μ on the model class M such that
{M, ,
SA 2 E
log K E
EM [
2
where T is the number of trajectories required by the algorithm Alg to achieve (s/12, 8) accuracy and A. ['] 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 c, h, a, the policy This." (Line 8) and the trusted transitions Tht (c) (Definition 3.1) can computed efficiently through linear programming. For simplicity, we fix context throughout this section and omit its dependence. We first note that if T for j s h 1 and n/h... computed, then The can be computed in poly(HSA) time. To this, for any (s, a, s'), we have by the definition of mhh.. that
Zh (8, TT)
Li (8, a; Th.s.a)
max SA + 7m 1
SA + 7m
rogm
(T
h,s,a
Then ss) E Th iff
Ti (8, OG mm
k.s,a
ph (s'(s,))))
SA + 7m reg
h,8,0 (Tm)
Smi
where the left-hand side can be computed given T for S h h- 1 poly(HSAA) time. Thus we only need to show how to compute nh.saa Assume want to compute - h.5, for h E [H], 3 E S, a A. Let fft (s, a) B,t.fa__,, [r^]] be the mean reward for model Mm-1 for
given Ttt for h < h
h, 8, Let Fm (F (s, a)).... be the vector of mean rewards for the model Mm-1. We consider following linear fractional program with the following decision variables:
-(at)) dh (d-1)
d/c
where d' (dnssammnssa d (dr.s.am (dnssam--)nssocc))sss.
We will use the variable dh to simulate the trusted occupancy measures and d to simulate the estimated occupancy measures from Mn-1 linear constraints. Concretely, consider the following

---

linear fractional program:
3,a.a.
max d
SA + 7m [11 m- (1 Fm-1))
-1 0,
Vh,8,a [H] x x88xxx
En d1..,,mmm112 = 1(s 81)
YsES,
Σ 3,0 d Ph
P 1- 1 s) =aa an
Vh, E [H] x S:
subject to
dha.amm 2 0,
Vh,8,a E [h] x xSXx x E
Σ dis.aam = L(s = = s¹);
E Exadeehh dn 1...m.. Ph- P E. dn..'.am
Vh S h E
This is linear fractional program of HSA+ ASA decision variables with HSA+ HS hSA hs constraints. It is clear from the linear constraints that this program simulates the MDPs Mm and Mm and the program obtains the right objective. Then, for this linear fractional program, we apply the Sharnes-Cooper transformation [7] transform into linear program. After the transformation. we can apply existing tools for solving linear programs (c.g., Lee and Sidford [23]) to solve for d. which encodes the policy h.3, in poly time.
25

---

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 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 your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] proovdddd 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 binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write 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 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.
is fine to include aspirational goals as motivation as long 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?
26

---

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 encouraged to create 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 few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-ted 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 faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 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 theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch 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
27

---

Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that 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 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 specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access 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 replicate the results, access 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 NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (c.g., a large language model), then there should either way to access this model for reproducing the results or 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 deseribe 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., 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 instruc- tions faithfully reproduce the main experimental results, 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 NeurlPS code and data submission guidelines (https://nips.co mbblcc/guudee//dddeuubisss PPolicyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not
28

---

including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results, See the NeurIPS code and data submission guidelines (https: Seuuidee/CodeSeee
wnips.cccpubll...
ss onnooicc) for more details.
The authors should provide instructions 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 subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if' applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
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 the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, 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, confi- dence intervals, statistical significance tests, 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, 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).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
29

---

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 96% CI, if the hypothesis of Normality of emrors 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 reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
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 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 into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https://neurii:
Answer: [Yes].
Justification: The research conducted in the paper conforms, every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations 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 theoretical work. There is no socictal 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 why the paper does not address societal impact.
30

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeet 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 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 unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition 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, 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 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 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 challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
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?
Question:
Answer: [NA].
Justification:
This paper does not existing assets.
Guidelines:
The answer NA means that the paper does not 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
31

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode com/dttasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
this information is not available online, the authors are encouraged to reach out to the asset's 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/mooel 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 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 research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According 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 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.
32

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing 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 the institution conducting the review.
33


Reviews:
[{"review_id": "dwPSTF4U44", "reviewer": "Reviewer_Bij6", "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.", "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?\n\n- Is there any lower bound on the number of calls for this problem?\n\n- Compared to contextual bandits, can the authors explain more about the technical challenges of CMDPs?", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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\n\n- This paper is purely theoretical. It would be interesting if the authors provide some experimental results to demonstrate the performance of the proposed algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">> This paper is purely theoretical. It would be interesting if the authors provide some experimental results to demonstrate the performance of the proposed algorithm. \"\"\"\n\nThis paper presents a purely theoretical result. We look forward to seeing experiments in future works.\n\n>> 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?\n\nThe value \\widehat{V}{m-1}^1(c) is the optimal value for dynamics \\widehat{P}{m-1}(c) and reward function \\widehat{R}{m-1}(c) which are known at the beginning of epoch m. Thus, it can be calculated through value iteration [1].\n\n[1] Markov decision processes: discrete stochastic dynamic programming\nML Puterman, 2014\n\n>> Is there any lower bound on the number of calls for this problem?\n\nThere is a lower bound on the switching cost of the scale $\\Omega(\\log \\log T)$ [2], where the switching cost is the number of switches in the learner’s randomized policy. A reasonable learner with an oracle would only switch its randomized policy after an oracle call. Thus, the number of oracle calls is larger than the number of switches, which implies a $\\Omega(\\log \\log T)$ lower bound on the number of oracle calls. We are happy to include a remark on this fact. However, to rigorously state this result requires more information-theoretic constraints and we leave such result to future works.\n\n[2] 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.\n\n\n>> Compared to contextual bandits, can the authors explain more about the technical challenges of CMDPs?\n\nThe main technical difficulty in general, for extending any result from contextual bandit to CMDPs is the unknown transition kernel. In fact, even if the reward function is given (which trivializes the contextual bandit problem since all rewards are known), the CMDPs are still hard to learn. Concretely, since the transition kernels are different in each layer, the learner has to ensure sufficient coverage to (nearly) ALL states-action pairs to estimate the corresponding transition kernel accurately. A more refined challenge in our setting is that if the learner only plans according to \\widehat{d}, which is the empirical estimation of the occupancy measure, the divergence between \\widehat{d} and the true occupancy measure will accumulate exponentially with respect to the horizon. To avoid such exponential blow-ups, we introduce the trusted transition in a layerwise fashion so that \\widetilde{d}, which is the trusted occupancy measure, is upper bounded by the true occupancy measure up to a constant.\n\nFor a detailed version, we refer to our general rebuttal.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tdxMF98Wsd", "reviewer": "Reviewer_eNC6", "summary": "The 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.", "questions": "See above.", "limitations": "N/a", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate your comments and have written a detailed technical overview as a general rebuttal. We believe it would be helpful for you to check before the detailed rebuttal.\n\n>> 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?\n\nYes, it can be implemented using ERM on the log loss. In fact, since ERM on log loss is MLE, we demonstrate in Example C.1 in appendix C that MLE is an offline oracle. This differs from Foster et al. and Levy et al. since they require online log loss guarantee.\n\n>> Why do you need to separate the data of each horizon for the model estimation?\n\nCould you clarify your question? I am reading two questions here: \n\nThe first one is, why not use all the data for model estimation, say at the end of each epoch? For this question, we separate the data so as to satisfy the i.i.d. input data assumption of the offline oracle. To perform the layerwise exploration-exploitation tradeoff,  the trajectory data generated in one epoch are NOT i.i.d. due to adaptively chosen policies in different segments. \n\nThe second one is why, for each layer h \\in [H], the estimator \\widehat{P}^h and \\widehat{R}^h are generated using separate data. This relates to the core challenge of our problem. If the learner only plans according to \\widehat{d}, which is the empirical estimation of the occupancy measure, the divergence between \\widehat{d} and the true occupancy measure will accumulate exponentially with respect to the horizon. To avoid such exponential blow-ups, we introduce the trusted transition in a layerwise fashion so that \\widetilde{d}, which is the trusted occupancy measure, is upper bounded by the true occupancy measure up to a constant. To ensure this property, at each layer, the model has to be estimated according to the policy that depends on the trusted occupancy measure from the last layer. \n\n>> Can you give some intuition regarding the definition of the trusted transitions?\n\nRecall the core problem mentioned above: we would like the trusted occupancy measure to be upper-bounded by the true occupancy measure up to a constant. For simplicity, assume the regret is always 0, then trusted transitions, according to our definition, are the ones that are most visited. Intuitively, such transitions are the ones that can be estimated most accurately. Indeed, such transitions can be estimated up to a multiplicative factor as we demonstrate in the proof of Lemma 3.4, which results in the trusted occupancy measure being upper-bounded by the true occupancy measure up to a constant.\n\n>> Related to this, can you explain how the policy cover is chosen?\n\nFor simplicity, assume regret is 0, then the policies in the policy cover maximize trusted occupancy measure which corresponds to exploration. When the regret is not 0, it balances the exploration-exploitation tradeoff.\n\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\nThanks for the suggestion. We will try to improve the structure of our paper.\n\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\nAt a very high level, one intuitive technical logic chain goes as follows: Since we would like to use an offline guarantee, we need to use the guarantee in the form d D_H(Phat, P). But the learner only knows \\hat{d} and can bound the divergence in value functions by \\hat{d} D_H(Phat, P). One immediate thought is to bound \\hat{d} D_H(Phat,P) by d D_H(Phat, P) + |d-\\hat{d}|. However, it will suffer from exponentially accumulating errors since |\\hat{d}-d| can only be written as summations of occupancy measures from the previous layers. So to avoid such an exponential explosion, we turn to multiplicative guarantees between d and \\tilde{d}.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fnM8uZn2nc", "reviewer": "Reviewer_Nj2f", "summary": "This 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.", "questions": "1. 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)?", "limitations": "The authors adequately address the limitations of their work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">> 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\nThank you for your valuable suggestions. Please refer to the general rebuttal for a technical overview of our methodology.\n\n>> 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\nThis paper is purely theoretical, but we hope to see empirical investigations in its possible extension to a journal version or future works.\n\n>> 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\nThe first to try would definitely be the MLE oracle as shown in Example C.1 in appendix C to enjoy theoretical guarantees. Concretely, since the MLE is indeed the ERM on log loss (i.e., cross-entropy loss), the first thing to try is probably ERM on cross-entropy loss.\n\n>> 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\nThe current technique needs to build the trusted occupancy measure, which unfortunately brings in the dependence of the state space S. For a problem with large state space, the guarantee is not ideal, but this could be alleviated through representation learning if the CMDP under consideration is intrinsically low dimensional, e.g. [1]. The dependence on the model class |M| is already logarithmic, which to our knowledge, is state-of-the-art. However, further structures in the model class might be helpful. We hope to explore these directions in future works.\n\n[1] Contextual Markov Decision Processes using Generalized Linear Models. Modi and Tewari. (2019).\n\n>> 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\nIt is not obvious how misspecification would affect the performance. On the positive side, LOLIPOP follows the stream of algorithms that adapt to misspecification, as seen in [2] and [3], which shows that IGW (which is a base component for LOLIPOP) is robust against misspecification. On the negative side, the proof technique requires the trusted occupancy measure to be multiplicatively upper-bounded by the occupancy measure. It might be non-trivial to maintain this guarantee with misspecification. Overall, we hope to see extension in this regard in the future.\n\n[2] Bypassing the monster: A faster and simpler optimal algorithm for contextual bandits under realizability\nD Simchi-Levi, Y Xu\n\n[3] Adapting to Misspecification in Contextual Bandits. \nDylan J. Foster, Claudio Gentile, Mehryar Mohri, Julian Zimmert\n\n>> 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\nAt a high level, the LOLIPOP algorithm promotes layerwise exploration-exploitation tradeoffs and planning with trusted transitions; it is quite conceivable to have implications in relevant topics. Specifically for constrained RL, one could imagine doing exploration-exploration tradeoffs in a constrained policy set. The LOLIPOP algorithm should have similar regret guarantees compared to the optimal policy in the constrained policy set.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "C9jXCtRo8p", "reviewer": "Reviewer_kcAb", "summary": "In 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$.", "questions": "1.\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.", "limitations": "1.\tThe regret bound is not tight in $H,S,A$.\n2.\tThe oracle might be inefficient.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.\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.", "weaknesses": "1.\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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">> In 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\nThanks for the suggestions!  We deferred the definition due to its length, but we will add an informal definition upfront for clarity. At a very high level, the intuition behind the multiplicative guarantee is as follows: when estimating a Bernoulli variable, if the empirical mean is greater than O(1/sample size), then with high probability, the empirical mean can be upper bounded by a constant multiple of the true mean.\n\n\n>> The 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\n\nWe respectfully disagree on this point. There are two significant advantages to using an oracle:\nWe would like to emphasize again the exponential difference between O(T) and O(\\log⁡T). With T known beforehand, the LOLIPOP algorithm achieves an oracle complexity of O(\\log⁡\\log⁡ T). We believe this is a substantial improvement since the gap can be double exponential.\nThe second advantage is the type of oracle required. Previously, to achieve near-optimal regret, algorithms required an online oracle. This is disadvantageous since online oracles are less straightforward to implement and necessitate updating the prediction for each round, which can be costly in practice. The offline oracles that the LOLIPOP algorithm uses can be implemented by ERM on the log loss. This is beneficial in practice since regression is much better understood than online updates, and the learner can choose to update only when a sufficient amount of new data has been gathered.\nOverall, we believe our algorithm is more practical.\n\n\n\n>> In 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). \n\nWe thank the author for sharing their vision. Continuing our argument, the MLE oracle is essentially ERM on log loss, making it more implementable and efficient than the online oracles typically considered in the literature. Additionally, due to the “online-to-batch” conversion, any online oracle can be converted to an offline one.\n\nWe are happy to discuss further the efficiency of offline oracles versus online ones. However, we note that this topic is somewhat outside the scope of this paper, as the efficiency and implementation of MLE is a classical statistical issue [5] rather than a decision-making one.\nRegarding linear functions, I am unsure what the reviewer means since we are focused on density estimation. The most popular class, logistic loss corresponding to the softmax distribution on states, is convex and thus allows for efficient implementation.\n\n[5] Maximum likelihood estimation: Logic and practice, SR Eliason, 1993\n\n>> Moreover, 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. I 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\nThe Hellinger distance has become quite standard in considerations related to MDP since Foster et al. (2021), as referenced by all the compared works. We are happy to add a citation to credit Foster et al. (2021). We chose the Hellinger distance for the same reason as the references: it can be effectively controlled through log-loss, specifically by MLE (log-loss regression), for our purpose.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
General rebuttal: 

We thank the reviewers for their positive reviews and will make sure to address the writing issues mentioned. Here, we clarify the technical challenges and our methodology. We will include a paragraph/section regarding this in the updated version of this paper.

Technical challenge and methodology:

The main technical difficulty in extending any result from contextual bandits to CMDPs is the unknown transition kernel. In fact, even if the reward function is given (which trivializes the contextual bandit problem since all rewards are known), CMDPs are still hard to learn. Concretely, since the transition kernels are different in each layer, the learner has to ensure sufficient coverage of (nearly) all state-action pairs to accurately estimate the corresponding transition kernel. 

A more refined challenge in our setting is that if the learner only plans according to $\widehat{d}$, which is the empirical estimation of the occupancy measure, the divergence between $\widehat{d}$ and the true occupancy measure will accumulate exponentially with respect to the horizon. The exponential blow-up is due to the following reason: since we would like to use an offline guarantee, we need to use the guarantee in the form $d D_H^2(\widehat{P}, P)$. However, the learner only knows $\widehat{d}$ and can bound the divergence in value functions by $\widehat{d} D_H^2(\widehat{P}, P)$. One immediate thought is to bound $\widehat{d} D_H(\widehat{P}, P)$ by $d D_H^2(\widehat{P}, P) + |d - \widehat{d}|$. However, it will suffer from exponentially accumulating errors since $|\widehat{d} - d|$ can only be written as summations of occupancy measures from the previous layers.

To avoid such an exponential explosion, we turn to multiplicative guarantees between $d$ and $\tilde{d}$. Concretely, we introduce the trusted transition in a layerwise fashion so that $\widetilde{d}$, which is the trusted occupancy measure, is upper bounded by the true occupancy measure up to a constant. To ensure this property, at each layer, the model has to be estimated according to the policy that depends on the trusted occupancy measure from the last layer. Intuitively, for simplicity, assume the regret is always zero; then trusted transitions, according to our definition, are the ones that are most visited. Such transitions are the ones that can be estimated most accurately up to a constant. At a very high level, when estimating a Bernoulli variable, if the empirical mean is greater than O(1/sample size), then with high probability, the empirical mean can be upper-bounded by a constant multiple of the true mean. Indeed, such transitions can be estimated up to a multiplicative factor, as we demonstrate in the proof of Lemma 3.4, which results in the trusted occupancy measure being upper-bounded by the true occupancy measure up to a constant. Meanwhile, to ensure sufficient coverage of all state-action pairs, we use the policy cover. For simplicity, assume regret is zero; then, the policies in the policy cover maximize the trusted occupancy measure, which corresponds to exploration. When the regret is not zero, the policy cover balances exploration and exploitation.

Meta-review: The paper explores a method that reduces the general Contextual Markov Decision Process (CMDP) problem to an offline density estimation task. The authors introduce an algorithm named LOLIPOP, which efficiently achieves near-optimal regret with a minimal number of oracle calls. The algorithm can also be applied to pure exploration tasks within reward-free reinforcement learning scenarios. The technical novelty of the paper is commendable which resolves an important open problem of learning CMDPs efficiently with offline oracles and rate optimal regret. Some empirical evaluation would have made the submission stronger though. 
Please incorporate all the reviewers' feedback in the final version.
Decision: Accept (poster)