The Benefits of Being Distributional: Small-Loss Bounds for Reinforcement Learning
Authors: Kaiwen Wang
Kevin Zhou
Runzhe Wu
Nathan Kallus
Wen Sun
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement Learning Theory
Distributional Reinforcement Learning
Small-Loss Bounds
First-order regret
TL;DR: We provide an explanation for the benefits of distributional RL through the lens of small-loss bounds, which scale with the instance-dependent optimal cost.

Abstract: While distributional reinforcement learning (DistRL) has been empirically effective, the question of when and why it is better than vanilla, non-distributional RL has remained unanswered.
This paper explains the benefits of DistRL through the lens of small-loss bounds, which are instance-dependent bounds that scale with optimal achievable cost.
Particularly, our bounds converge much faster than those from non-distributional approaches if the optimal cost is small.
As warmup, we propose a distributional contextual bandit (DistCB) algorithm, which we show enjoys small-loss regret bounds and empirically outperforms the state-of-the-art on three real-world tasks.
In online RL, we propose a DistRL algorithm that constructs confidence sets using maximum likelihood estimation. We prove that our algorithm enjoys novel small-loss PAC bounds in low-rank MDPs.
As part of our analysis, we introduce the $\ell_1$ distributional eluder dimension which may be of independent interest. 
Then, in offline RL, we show that pessimistic DistRL enjoys small-loss PAC bounds that are novel to the offline setting and are more robust to bad single-policy coverage.

The Benefits of Being Distributional: Small-Loss Bounds for Reinforcement Learning
Kaiwen Wang
Kevin Zhou
Runzhe Wu Nathan Kallus
Wen Sun
Cornell University
(kw437, 1,,2k,12,,,,,,, rw646, kallus, ws455) corrsell..o..el.c.ooc
Abstract
While distributional reinforcement learning (DistRL) has been empirically effective, the question of when and why it is better than vanilla, non-distributionll RL has remained unanswered. This paper explains the benefits of DistRL through the lens of small-loss bounds, which are instance-dependent bounds that scale with optimal achievable cost. Particularly, our bounds converge much faster than those from non-distributional approaches if the optimal cost is small. As warmup, we propose distributional contextual bandit (DistCB) algorithm, which we show enjoys small-loss regret bounds and empirically outperforms the state-of-the-art on three real-world tasks. In online RL, we propose DistRL algorithm that constructs confidence sets using maximum likelihood estimation. We prove that our algorithm enjoys novel small-loss PAC bounds in low-rank MDPs. As part of our analysis, we introduce the & distributional eluder dimension which may be of independent interest. Then, in offline RI,, we show that pessimistic DistRL enjoys small-loss PAC bounds that are novel to the offline setting and are more robust to bad single-policy coverage.
Introduction
The goal reinforcement learning (RL) is to learn a policy that mininizes/maximizes the mean loss/return (i.e., cumulative costs/rewards) along its trajectory. Classical approaches, such as Q- learning [Mnih et al., 2015] and policy gradients [Kakade, 2001], often learn Q-functions via least square regression, which represent the mean loss-to-go and act greedily with respect these estimates. By Bellman's equation, Q-functions suffice for optimal decision- maasiinnmmaiingg and indeed these approaches have vanishing regret bounds, suggesting only need to learn means well [Sutton and Barto, 2018]. Since the seminal work of Bellemare et al. [2017], however, numerous developments showed that learning the whole loss distribution stratospheric balloon navigation [Bellemare et al., 2020], robotic grasping [Bodnar et al., 2020], algorithm discovery [Fawzi al., 2022] and game playing benchmarks [Hessel et al., 2018, Dabney et al., 201 8a, Barth-Maron et al., 2018J. In both online [Yang et al., 2019] and offline RL [Ma al., 2021], distributional RL (DistRL) algorithms often perform better and use fewer samples in challenging tasks when compared to standard approaches that directly estimate the mean.
actually yield tfate-of-the-an performance in
Despite learning the whole loss distribution, DistRL algorithms only the mean of the learned distribution for decision making, not extracting any additional information such as higher moments. In other words, DistRL is simply employing different and seemingly roundabout way of learning the mean: first, learn the loss-to-go distribution via distributional Bellman equations, and then, compute the mean of the learned distribution. Lyle al. [2019] provided some empirical explanations of the benefits of this two-step approach, showing that learning the distribution, e.g., its moments or quantiles, is auxiliary task that leads to better representation learning. However, the theoretical
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

question remains: does DistRL, i.e., learning the distribution and then computing the mean, yield provably stronger sinite-sample guarantees and if so stronger how and when?
In this paper, we provide the first mathematical hasis for the benefits of DistRI. via the lens of small-loss bounds, which are eestance-depenenn bounds that depend on the minimum achievable cost in the problem [Agarwal et al., 2017].' For example in linear MDPs, typical worst-case regret bounds scale on the order of poly(d, H)vK, K, where d is the feature dimension, H is the horizon, and K is the number of episodes [Jin et al., 2020b]. In contrast, small-loss bounds will scale on the order of poly (d, H. V V + poly(d, H) log(K)) where V min, V is the optimal expected cumulative cost for the problem. We assume cumulative costs are normalized in [0, 1] without loss of generality. As V* becomes negligible (approaches 0), the first term vanishes and the small-loss bound yields faster convergence rate of (poly(d, H) log (K), compared to the Oppoly (d, H)V K)) rate in standard uniform bounds. Since we always have V* I 1, small-loss bounds simply match the standard uniform bounds in the worst case.
As warm-up, we show that maximum likelihood estimation (MLE), Le., maximizing log-likelihodd, can be used obtain small-loss regret bounds for contextual bandits (CB), i.e., the one-step RL setting. Then, we turn to the online RL setting, and propose an optimistic DistRL. algorithm that optimizes over confidence sets constructed via MLE applied the distributional Bellman equations. We prove our algorithm attains the first small-loss PAC bounds in low-rank MDPs [Agarwal et al., 2020]. Our proof uses novel regret decomposition with triangular discrimination and also introduces the { distributional eluder dimension, which generalizes the & distributional eluder dimension of Jin al. [2021a] and may be of independent interest. Furthermore, we design an offline distributional RL algorithm using the principle of pessimism, and show our algorithm obtains the first small- bounds in offline RI. Our offline small-loss bound holds under the weak single-policy coverage. Notably, our result has a novel robustness property that allows our algorithm to strongly compete with policies that either are well-covered or have loss, while prior approaches solely depended on the former. Finally, we find that our distributional CB algorithm empirically outperforms existing approaches in three challenging CB tasks.
Our key contributions are as follows:
1. As warm-up. we propose a distributional CB algorithm and prove that it obtains a small- loss regret bound (Section 4). We empirically demonstrate it outperforms state-of-the-arr CB algorithms in three challenging benchmark tasks (Section 7).
2. We propose distributional online RL algorithm that enjoys small-loss bounds in settings with low l distributional cluder dimension, which we show can always capture low-rank MDPs. The & distributional cluder dimension may be of independent interest (Section 5).
3. We propose a distributional offline RL algorithm and prove that it obtains the first small-loss bounds in the offline setting. Our small-loss guarantee exhibits a novel robustness to bad coverage, which implies strong improvement over more policies than existing results in the literature (Section 6).
In sum, we show that DistRL can yield small-loss bounds in both online and offline RL, which provide concrete theoretical justification for the benefits of distribution learning in decision making,
2 Related Works
Theory of Distributional RL Rowland et al. [2018, 2023) proved asymptotic convergence guar- antees of popular distributional RL algorithms such as C51 [Bellemare et al., 2017] and QR-DQN [Dabney et al., 2018b]. However, these asymptotic results do not explain the benefits of distributional RL over standard approaches, since they do not imply stronger finite-sample guarantees than those obtainable with con-distributional algorithms. In contrast, our work shows that distributional RL yields adaptive unite-sample bounds that converge faster when the optimal cost of the problem is small. Wu al. [2023] recently derived finite- sample bounds for distributional off- oof-colicy evaluation with MLE, while our offline RL section focuses on off- off-policy optimization.
order" generally refers to bounds that scale with the optimal value, either the maximum reward or the minimum cost. To highlight that minimizing cost, we call our bounds "small-loss".

---

First-order bounds in bandits When maximizing rewards, first-order remall-retumm" bounds can be easily derived from EXP4 [Auer et al., 2002], since receiving the worst reward with probability (w.p.) ò contributes at most R* to the regret². When minimizing costs, receiving the worst loss 1 w.p. may induce large regret relative to I. if T.* is small. To illustrate, R then all policies are optimal, so no learning is needed and the small-return bound is vacuous. Yet if I, = 0, sub-optimal policies may have a large gap from L", so small-loss bounds in this regime are meaningful. Small-loss bounds are achievable in multi-arm bandits [Foster et al., 2016], semi-bandits [Neu, 2015, Lykouris et al., 2022], and CBs [Allen-Zhu et al., 2018, Foster and Krishnamurthy, 2021J.
First-order bounds in RL Jin et al. [2020a], Wagenmaker et al. [2022] obtained small-return regret for tabular and linear MDPs via concentration bounds that scale with the variance. The idea is that the return's variance is bounded by some multiple of the expected value, which is bounded by V* in the reward- maximizing setting, i.e., Var(En 2h = c. V* last inequality fails in the loss-minimizing setting, so the variance approach does not easily yield small-loss bounds. Small-loss regret for tabular MDPs was resolved by Iee et al. [2020, Theorem 4.1] using online mirror descent with the log-barrier on the occupancy measure. Moreover, Kakade et al. [2020, Theorem 3.8] obtains small-loss regret for linear-quadratic regulators (LQRs), but their Assumption 3 posits that the coefficient of variation for the cumulative costs is bounded, which is false in general even in tabular MDPs. To the best of our knowledge, there are no known first-order bounds for low-rank MDPs or in offline RL.
I c. V*. However, the
Risk-sensitive RL A well-motivated use-case of DistRL is risk-sensitive RL, where the goal is to learn risk-sensitive policies that optimize some risk measure, e.g., Conditional Value-at-Rin (CVaR), of the loss [Dabney et al., 20] Orthogonal to risk-sensitive RL, this work focuses on the benefits of DistRL for standard risk-neutral RL. Our insights may lead to first-order bounds for risk-sensitive RL, which we leave as future work.
3 Preliminaries
As warmup, we begin with the contextual bandit problem with an arbitrary context space X, finite action space A with size A and conditional cost distributions C: X' x A - A[O,, 1]). Throughout, we fix some dominating measure on [0, 1] (e.g.. Lebesgue for continuous or counting for discrete) and let ( IJ) be all distributions on [0. 1] that are absolutely continuous with respect to A. We identify such distribution with its density with respect to X, and we also write Cly x, a) for (C(z, (v). Let K denote the number of episodes. At each episode l E [K] the learner observes a context Zy E X, samples an action 0g E AA and then receives a cost Ct ~ C(2,, at), which we assume to be normalized, i.e., Ct E [0, 1]. The goal is to design learner that attains low regret with high probability, where regret is defined as
Regretes (K) ExE C(zk, O(IK, (Ix)),
where f J for any f A(O 1]) and x (**x) arg min C)z,,
The focus of this paper is reinforcement learning (RL) under the Markov Decision Process (MDP) model, with observation space X, finite action space A with size A, horizon H, transition kernels PA, X' x A + (Ax) and cost distributions C₁ X' x A + A ([0. 1l) at each step h E [H]. We start with the Online RL setting, which proceeds over K episodes as follows: each episode & E [K], the learner plays policy n² E IX A(A)]"; we start from a fixed initial state T1; then for each h = 1, 1,2, ... .., H, the policy samples an action a₁ ~ * (2x), receives a cost Cr ~ (ar,,, an). and transitions to the next state F441 Pr (m an.). Our goal is to compete with the optimal policy that " minimizes expected the loss, i.e., n E arg min V where V Ex Σ Σm= ch. Regret bounds aim to control the learner's regret with high probability, where regret is defined as,
Regrel RI. (K) = Σx= V* V*.
If the algorithm returns a single policy म. it is desirable to obtain a Probably Approximately Correct (PAC) bound on the sub-optimality of R, i.e., v V*
Assume rewards/losses in [0, 1] and R¹ /L" is the maximum/minimmun expected reward/loss.
3

---

The third setting we study is Offfine RL, where instead of needing to actively explore and collect data ourselves, we are given H datasets D1 D2.... ........ DH to learn a good policy म. Each Di, contains N i.dd. samples (Ih,5; On,i; Ch, Kh,s)) from the process (Ih,i (h,s) Vh, Ch,s ~ Ch (a Oh,i), *h,, Pa (1 (h,i), where M/ E A(X A) is arbitrary, e.g. the visitations of many policies from the current production system. The goal is to design an offline procedure with PAC guarantee for म. which should improve over the data generating process.
Distributional RL For policy a and h E [H], let (1 an) E A( 1)) denote the dis- tribution of the loss-to-go Σish Q conditioned on rolling in π from The ag. The expectation of the above is (T), an)) = (1,, an) and V (2x) = Ea (ax.) (4, 1Q an)]. We use 27, Q% V to denote these quantities with n². Recall the regular Bellman operator acts on function f: X' x A + [0, 1] as follows: Th ((x,a)= Cz, a) + Ey (a,a),'' (xx))) a')]- Analogously, the distributional Bellman operator [Morimura et al., 2012, Bellemare et al., 2017] acts on conditional distribution d x A(0,, 1)) follows: Tr.:d((a d(x, a) & Cn d(x' a'),
where x' (r. a) d' ~ (xx') and 2 denotes equality of distributions. Another way think about the distributional Bellman operator is that a sample 2 T. d(xx a) is generated as fol- low: == + y, where e Ca(r, a), x' ~ Pa(z, ~ m<=²),)) V. "ila',a') We will also use the Bellman optimality operator TØ and its distributional variant Thi defined follows: d')] and Th d(x,a) a) 2 Cz,z,a) d(x', where x' ~ Pr(x, a), arg min, d(('' a). Please see Table for an index of notations.
Ty f(r, a) (z, + (m..)
4 Warm up: Small-Loss Regret for Distributional Contextual Bandits
In this section, we propose an efficient reduction from CB online maximum likelihood estimation (MLE), which is the standard tool for distribution learning that we will use throughout the paper. In our CB algorithm, we balance exploration and exploitation with the reweighted inverse gap weighting (ReIGW) of Foster and Krishnamurthy [2021], which defines distribution over actions given predictions f E A and parameter E IR++: setting b arg min f(a) as the best action with respect to the predictions, the weight for any other action a + b is,
f(b)
ReIGW (f.)a) == Af(b) + (f((a) - f(b))'
and the rest of the weight is allocated to b: ReIGW, (F. = -1 -ooy ReIGW (f,
Algorithm Distributional CB (DISTCB)
1: Input: number of episodes K, failure probability 8, ReIGW learning rate 7.
2: Initialize any cost distribution (()
., 4: Observe context Zk-
3: for episode k = 1,2, ..., K do
(F()))) 6: Observe cost C ~ C'I'k ax) and update online MLE oracle with ((Ik, ax), cx).
5: Sample action ag pk ReIGW(F(() (Ik, 7) from Eq. (1).
7:
end for
We propose Distributional Contextual Bandit (DistCB) in Algorithm a two-step procedure for each episode & [K]. Upon seeing context æv. DISTCB first samples an action a from ReIGW generated
by means of our estimated cost distributions for each action, i.e., f(a) F(() (Ixk, a), Va E A (Line 5). Then, DISTCB updates (xx 2g, a&) by maximizing the log likeeliodoood to estimate the conditional cost distribution C' Fy. a4)) (Line 6). Formally, this second step is achieved via an online MLE oracle with realizable distribution class FCB C X x A + A(O0, 1)); let Regret (K) be some upper bound on the log- likeelhodoood regret for all possibly adaptive sequences {Øx, ak, Ck [
k= log C(c& Ik. ax) log (b) (cx Ik, ax) S Regret (K). (K).
Under realizability, C E FCB, we expect Regret (K) E (log((K), For instance, if Fon is finite, exponentially weighted average forecaster guarantees Regret log (K) S log |Fa| Diesa-Bianchi and Lugosi, 2006, Chapter 9]. We now state our main result for DISTCB.

---

Theorem 4.1. For any ó E (0, w.p. at least 1 6, running DISTCB with y = 10.4 V
V 112 40A(C" 1/8) has regret scaling with C' = Ek=1 minge C(r,,a),
Regret ( (K) - v AC* Regret (K) log(1/5) 2300 (Regretho (K) + log(1/8))
The dominant term scales with the optimal sum of costs v C which shows that DISTCB obtains small-loss regret. DISTCB is also computationally efficient since each episode simply requires computing the ReIGW. FastCB is the only other computationally efficient CB algorithm with small- loss regret [Foster and Krishnamurthy, 2021, Theorem 1J. Our bound matches that of FastCB in terms of dependence on A, C* and log(1/6), Our key difference with FastCB is the online supervised learning oracle: in DISTČB, we aim to learn the conditional cost distribution by maximizing log- likelihood, while FastCB aims to perform regression with the binary cross- ross-entropy loss. In Section 7, we find that DISTCB empirically outperforms SquareCB and FastCB three challenging CB tasks, which reinforces the practical benefits of distribution learning in CB setting.
4.1 Proof Sketch
First, apply the per-round inequality for ReIGW [Foster and Krishnamurthy, 2021, Theorem 4] to get,
K
A C(ks ax) + (F() as) Osk,, 27
Regret (K) 2
Eag ~Pk
f(h ax) + C(sp; ax)
k=1
For any distributions S,g A(0,, 1J), their triangular discrimination' is defined as D^ 9) J ((())=())) dA(y). The key insight is that * can be bounded by the triangular discrimination fout
of ((K) ax) and C(sk: ax): by Cauchy-Schwarrd Schwartz and y² < y for E [o, 1]. we have f - 9 = [(((()) ((u))dM()) S V v(f(() + alo))dM()) S (f())=s())) dA(y), and hence, f())9(())
|f 1 S (f+ g).
(41)
So, Eq. (A1) implies that * (up universal constants) to the squared Hellinger distance, Foster al. [2021, Lemma A. 14] implies the above can be bounded by the online MLE regret, so w.p. at least - 5, we have
is bounded by (f (sk, ax) || C(sk, ax)). Since Da is equivalent
Regref (K) - Σ (C(sk, ax) + log(1/8)) + > (Regret (K) + log(1/8)
K
From here, we just need to rearrange terms and set the correct Appendix C contains the full proof.
5 Small-Loss Bounds for Online Distributional RL
We now extend our insights to the online RL setting and propose a DistRL perspective on GOLF [Jin et al., 2021a]. While GOLF constructs confidence sets of near-minimizers of the squared Bellman error loss, we propose to construct these confidence sets using near-maximizers of the log-likelihood loss approximate MLE. To leverage function approximation for learning conditional distributions, we use a generic function class F S (X' x A - A[0,, 1))) " where each element f E F is a tuple (fi,..., fH) such that each fr. is candidate estimator for Zj, the distribution of loss-to-go El & under -*. For notation, fH+1 (z, a) = åg denotes the dirac at zero for all x, a.
We now present our Optimistic Distributional Confidence set Optimization (O-DISCO) algorithm in Algorithm 2, consisting of three key steps per episode. At episode & E [K]. O-DISCO first identifies the (8) with the minimal expected value at h 1 over the previous confidence set FR-1 (Line 4). This step induces global optimism. Then, O-DISCO collects data for this episode by rolling in with the greedy policy -² with respect to the mean of ((k) (Line 6). Finally, O-DISCO constructs a
Triangular discrimination is also known Vincze-Le Cam divergence [Vineze, 1981, Le Cam, 2012).
5

---

Algorithm 2 Optimistic Distributional Confidence set Optimization (O-DISCO)
1: Input: number of episodes K, distribution class F, threshold B.
2: Initialize Dv.0 - 0 for all h E [H] and set Fu F.
3: for episode k 1, ..... K do
4: Set optimistic estimate f(k) arg min SEFL-, ming f1(
5: Set (x) arg min, f (z.a).
6: Roll out œ and obtain trajectory 11,k, d1,,,
OHH, CH.k- For each h E [H) augment the dataset D,, = Di,K-1 ( (T 4h,ks Ch,k, 1.)]
7: For all f) E [M] x F, sample v ~ Tnti(dh, d') and a' = arg min, fn+ (2 a)
where (Fh,i; dh,i; Ch.s; .),,)) .) the i-th datapoint of D Then, set 2h. = Ch + Vn and define the confidence set
Fx 32=SEEF log fn This an.)) 2
max ge
logg(z/ Thas an.i) - 78, vh E [M]]
8: end for
9: Output: के 1.
confidence set F& by including a function / if it exceeds a threshold on the log-likelihood objective for all steps h simultaneously (Line 7). This step is called local fiting as each f E F& has the property that fa is close-in-distributio to Th fa+1 for all h. We highlight that O-DISCO only learns the distribution for estimating the mean, i.e, Lines and 6 only use the mean f. This seemingly roundabout way of estimating the mean is exactly how distributional RL algorithms such as C51 differ from the classic DQN.
using data <//,; ~ Thi (IThh (n,z)
To ensure that MLE succeeds for the Cimporal-Diiference (TD) style confidence sets, we need the following distributional Bellman Completeness (BC) condition introduced in Wu et al. [2023].
Assumption 5.1 (Bellman Completeness). For all n, h E [H], fA41 FR+1 =
T,. TTr.Dhntt sh+11
Fb.
5.1 The & Distributional Eluder Dimension
We now introduce the l distributional eluder dimension. Let S be an abstract input space, let V be set of functions mapping S + R and let 'D be a set of distributions on S.
Definition 5.2 dir-distributtinnal eluder dimension). For any function class V C S - R. distribution class D S (88) and 0, the 4y f,-distributional eluder dimension (denoted by DE,, D,,)) is the
length L of the longest sequence d(!), dd²), (I.) D such that there exists e 2 E, such that for all / [7.), we have that there exists V such that [E fl > and also Σ Eggo S1P S SP.
When p = 2, this is exactly the lz distributional eluder of Jin et al. [2021a, Definition 7]. We're particularly interested in the p = case, which can be used with MLE's generalization bounds. The following is a key pigeonhole principle for the l distributional cluder dimension.
Theorem 5.3. Let C supdED, [Eafl [Euf] be the envelope. Fix any K N and sequences (a), f(K) C V, ((1))...,,d(k)) (K()) C D. Let B be a constant such that for all k E [K], we have,
E Egts [(]] < 3. Then, for all ls E [K], we have
Σ
Ego /00/ < inf oce51
( (log(C//)) + ke).
As we'l II see later, Theorem 5.3 the key tool that transfers triangular discrimination guarantees on the training distribution to any test distribution. Another key property is that the dimension generalizes the original & dimension Jin et al. [2021a].
Lemma 5.4. For any V, D and
0, we have DE, (V, D, DF2(W D..).
Finally, we note that our distributional eluder dimension generalize the regular & eluder from Liu et al. [2022], which can be seen by taking D to be dirac distributions.

---

5.2 Small-Loss Bounds for O-DISCO
We will soon prove small-loss regret bounds with the "Q-type" dimension, where "Q-type" refers to the fact that S = X x A. While low-rank MDPs are not captured by the "Q-type" dimension, they are captured by the "V-type" dimension where S = X [Jin et al., 2021a, Du et al., 2021]. For PAC bounds with the V-type dimension, we need to slightly modify the data collection process in Line 6 with uniform action exploration (UAE). Instead of executing x² for a single trajectory, partially roll-out Th for H. times where for each h E [H], we collect Ih, ~ di take a random action an.k (A) observe Ch.K (T (n, ~ Pa. (I a),s) and augment the dataset DR,A= = P (Th,kk ahiki Ch,ss w)) The modified algorithm is detailed in Appendix B.
We lastly need to define the function and distribution classes measured by the distributional eluder dimension. The Q-type classes are Dn = -(1x, a) + (2x a) : IT E II} and Vn. {(x,a) + Da(S(x; a) || T. [(x a)) : E Similarly, the V-type classes are Dr.y fo + d7) : E I and 1n = fIr + En-Unitt (D) a) || TT. (T,.)(z,,)) a)]] nally, define (e) = max, DE, (Vn, e) and (E) = maxa DE, (Wn,:: D,,,,))
F]. Fi-
Theorem 5.5. Suppose DistBC holds (Assumption 5.1). For any 8 E (0, 1), w.p. at least 1 8. running O-DISCO with B = log(H K (FF/8) guarantees the following regret bound.
Regret D (K) I 160H / KV* (L/K)) log(K) + 18000 DE, (L/K) log B.. UFAAE TRUE (Algorithm 4), then the learned mixture policy ā is guaranteed satisff,
AV* 1 (1/K) K) log(K)B 18000 o log K)B
v V* < 160H
K
K
Compared prior bounds for GOLF [Jin et al., 2021a], the leading vK terms in our bounds enjoy the same sharp dependence in H, K and the eluder dimension. Our bounds further enjoy one key improvement: the leading terms are multiplied with the instance-dependent optimal cost V* giving our bounds the small- loss- property. For example, V S O(1/VK), then our regret bound converges at a fast O DE, (L/K) log K)) 3)) rate. While there are existing first-order bounds in online RI., our bound significantly improves on their generality. For example, Zanette and Brunskill [2019]. Jin et al. [2020a], Wagenmaker et al. [2022] used Bernstein bonuses that scale with the conditional variance and showed that careful analysis can lead to "small-return" bounds in tabular and linear MDPs. However, "small-return" bounds do not imply "small-loss" bounds and "small-loss" bounds are often harder to obtain'. While it is possible that surgical analysis with variance bonuses can lead to small-loss bounds in tabular and linear MDPs, this approach may not scale to settings with non-linear function approximation such as low-rank MDPs.
On Bellman Completeness Exponential error amplification can occur in online and offline RL under only realizability of Q functions [Wang et al., 2021a,b,c, Foster et al., 2022]. With only realizability, basic algorithms such as TD and fattedQ--Evalueee to bad fixed point solutions [Tsitsiklis and Van Roy, 1996, Munos and Szepesvári, 2008, Kolter, 201 I]. As result, BC has risen as a de facto sufficient condition for sample efficient RL [Chang et al., 2022, Xie et al., 2021, Zanette et al., 2021]. Finally, we highlight that our method can be easily extended to hold under generalized completeness, i.e., there exist function classes Gn such that fh+1 E Fh+1 = Ty.DJn++ E 9n [as in Jin et al., 2021a, a, Assumption 14). Simply replace maXger, in the confidence set construction with maxyes.. While adding functions to F may break BC (as BC is not monotonic),
(FQE) can diverge or converge
always augment G to satisfy generalized completeness.
Computational complexity When taken as is, OLIVE [Jiang et al., 2017], GOLF, and our al- gorithms are version space methods that suffer from computational drawback: optimizing over the confidence set is NP-hard [Dann et al., 2018]. However, the confidence set is purely for deep exploration via optimism and can be replaced by other computationally efficient exploration strategies. For example, s-greedy suffices in problems that don't 't require deep and strategic exploration, i.e., a large myopic exploration gap [Dann et al., 2022]. With e- e-greedy, replay buffer, and discretization, our algorithm essentially recovers C51 [Bellemare et al., 2017]. We leave developing and analyzing computationally efficient algorithms based on insights promising future work.
In Appendix J. we show slight modification of our approach also yields "small-retum"" bounds.
7

---

5.3 Instantiation with Low-Rank MDPs
The low-rank MDP [Agarwal et al., 2020] is a standard abstraction for non-linear function approx- imation used in theory [Uehara al., 2021] and practice [Zhang al., 2022, Chang al., 2022].
Definition 5.6 (Low-rank MDP). A transition model Pa : X x A -> (X)) has rank d if there exist unknown features 0% : X' x A + Rd 1% A' + IR such that (a 2, a) = (1 (x') for all x." Also, assume max 1 (x. a) 22 S 1 and II gdyt ll < ||ll| vd for all functions g: + R. The MDP is called low-rank if Pa is low-rank for all h E [M].
We now specialize Theorem 5.5 low-rank MDPs with three key steps. First, we bound the V-type eluder dimension by (E) I ((d(o/(//))) which is known result that we reproduce in Theorem G.4. The next step requires access to a realizable & class, i.e. for all h E [H], 0% E Φ, which is standard assumption for low-rank MDPs [Agarwal et al., 2020, Uehara et al., 2021, Mhammedi et al., 2023]. Given the realizable Q, we can construct a specialized F for the low-rank MDP: Flln Flin ..... Flin x Flin where Flin H+1 {ōo} and for all h [H].
H+1
Fltn { 5( x, a) (0(x; (), w(z))
e E w w [0,1]
Rd
(2)
s.t. max | avd and max (ő(z, a). mx))) {).
8,0,2
where maX ( (z) | a. is the largest mass for the cost-to-go distributions. In Appendix D, show that Flll satisfies DistBC. Further, if costs are discretized into uniform grid of M points, its bracketing entropy is bounded by O(dM log (b). Discretization is necessary to bound the statistical complexity of Flin and is common in practice, e.g., C51 and Rainbow both set M = 51 which works well in Atari games [Bellemare et al., 2017, Hessel et al., 2018].
Theorem 5.7. Suppose the MDP is low-rank. For any 5 E (0,1), 1). w.p. at least - 6, rumning O-DISCO with UAE=TRUE and with Flin as described above learns a policy म such that,
AdV (dM + log log(11)))ggg +1/8))
2dH? (dM + loog( (og(((1/8))
V$ - V* E O H
K
K
Proof. As described above, we have (1/K) S Oddoog(d)) and B log( HK/5) dM log (bb Since DistBC satisfied by Flin plugging into Theorem 5.5 gives the result.
This is the first small-loss bound for low-rank MDPs, and for online RL with non-linear function approximation in general. Again when V* v Ö(1/K), O-DISCO has fast (1/K) convergence rate which improves over all prior results that converge at slow ((1/V/) rate [Uehara al., 2021].
5.4 Proof Sketch of Theorem 5.5
By DistBC (Assumption 5.1), we can deduce two facts about the construction of F&: (1) Z* E Fr, and (ii) elements of Fs almost satisfy the distributional Bellman equation, i.e., for all [H), we have
1 E,, [ az)) < ((3) where 5n, (Z1, an) = Da(ft (zn, an) Il T Jh41 (#h, az)). Next, we derive a corollary of Eq. (A1):
(R)
f f- gl S V 4g (a(f D DD(f
(A2)
To see why this is true, apply AM-GM to Eq. (A1) to get 2(F g) S f Dz || g), which simplifies f < 3g + D&(f g). Plugging this back into Eq. (A1) yields Eq. (D2). Then, by iterating Eq. (Az) and AM-GM, we derive a self- eeff-bounding lemma: for any f, क. h, we have f441(2, az))] (Lemma H.3).
fa on) S (2h, an) + H H" Eg Bn on (f( at) T
8

---

h Since Th" Jut1 (x, a) = -Th" ( (1 (2 and Th T T Shh = Th shtt we have
Va - V S V= nー - 51 (1I,
(optimism from fact (i))
- Bnp Thi p(k)
Jut (an, an) ( an)
(performance difference)
5 <2Σ V /Ba 15 (In, and + ( an B ( an)l
(Eq. (42))
2 / + H h== B 18 and HE (O, an)l.
E
(Lemma H.3)
The implicit inequality V= - V* S V + H Ex (On, (Zh, an) HE,, [Ōn, (Tp, as can then be obtained by AM-GM and rearranging. The final step is to sum over k and bound Ek E, (Ta, as)] via the eluder dimension's pigeonhole principle (Theorem 5.3 applied with fact (ii)). Please see Appendix for the full proof.
Small-Loss Bounds for Offline Distributional RL
We now propose Pessimistic Distributional Confidence set Optimization (P-DISCO) Algorithm 3), which adapts the distributional confidence set technique from the previous section the offline setting by leveraging pessimism instead of optimism. Notably, P-DISCO is simple two-step algorithm that achieves the first small-loss PAC bounds in offline RL. First, construct a distributional confidence set for each policy based on similar log- likelihood thresholding procedure as in O-DISCO, where the difference is we now use data sampled from Th D fh+1 instead of Thr Sa+1- Next, output the policy with the most pessimistic mean amongst all the confidence sets.
Algorithm 3 Pessimistic Distributional Confidence set Optimization (P-DISCO)
1: Input: datasets D D..... DHr distribution function class F, threshold 3, policy class IT. 2: For all (h, TT) E [#] F x II, sample y fn-1 (141((%hh)) where (Kh,s Oh,i, Ch,a, sk,s) is the i-th datapoint of Dn. Then, set x Ch, vhi and define the confidence set,
N
Fr SEF: Dog fn (1/ Fhis an.a) 2 max
log g g( for Thii an.i,an)) - 73, wh E (H)
gEFi
3: For each II, define the pessimistic estimate f arg max max Bawa(s)) B (((
4: Output: के = arg max (awn(1)) (x1,
In offline RL, many works made strong all-policy coverage assumptions [Antos al., 2008, Chen and Jiang, 2019]. Recent advancements [Kidambi et al., 2020, Xie et al., 2021, Uchara and Sun, 2022, Rashidinejad et al., 2021, Jin et al., 2021b] have pursued best effort guarantees that aim to compete with any covered policy H. with sub-optimality of the learned रे degrading gracefully as coverage worsens. The coverage measured by the single-policy concentrability C*² max, [ddi/aon lbo We adopt this framework and obtain the first small-loss PAC bound in offline RL.
Theorem 6.1 (Small-Loss PAC bound for P-DISCO), Assume Assumption 5.1. For any 5 E (0. 1), w.p. at least 1 - 5, rumning P-DISCO with 3 = log H ())) learns a policy स that enjoys the following PAC bound with respect to any comparator policy II:
/
CRVVB N
30H2CB8
- S 99H
N
To the best of our knowledge, this is the first small-loss bound for offline RI, which we highlight illustrates a novel robustness property against bad coverage. Namely, the dominant term not only scales with the coverage coefficient C but also the comparator policy's value V In particular, P-DISCO can strongly compete with comparator policy स if one of the following is true: (i) has good coverage over मं so the (11vVN) term is manageable; we may even obtain a fast ( rate. Thus, P-DISCO has two chances at strongly competing with म. while conventional offline RL methods solely rely on (i) to be true.
(ii) गे has small-loss, in which case

---

7 Distributional CB Experiments
Algorithm:
SquareCB
FastCB
DistCB (Ours)
We now compare DISTCB with SquareCB [Foster and Rakhlin, 2020] and the state-of- the-art art CB method FastCB [Foster and Krish- namurthy, 2021], which respectively minimize the squared loss and log loss for estimating the conditional mean. The key question we investigate here whether learning the con- ditional mcan via distribution learning with MLE will demonstrate empirical benefit over the non-distributional approaches. We consider three challenging tasks that are all derived from real-world datasets and we briefly describe the construction below.
King County Housing [Vanschoren et al., 2013]
All episodes .756 (.0007). .734 (.0007) .726 (.0003)
Last 100 ep. .725 (.0012) .719 (.0013) .708 (.0019)
Prudential Life Insurance [Montoya et al., 2015]
All episodes .456 (.0082) .491 (.0029) .411 (.0038) Last 100 cp. .481 (.0185) .474 (.0111) .388 (.0086)
CIFAR-100 [Krizhevsky, 2009]
All episodes 8872 (.0010) .856 (.0016) .838 (.0021)
Last 100 ep. .828 (.0024). .793 (.0031) .775 (.0027)
Table 1: Avg cost over all episodes and last 100 episodes (lower is better). We report 'mean (sem)' over 10 seeds.
King County Housing This dataset consists of home features and prices, which we normal- ize to be in (0, 1]. The action space is 100 evenly spaced prices between 0.01 and 1.0. If the learner overpredicts the true price, the cost 1.0. Else, the cost is 1.0 minus predicted price.
Prudential Life Insurance This dataset contains customer features and an integer risk level in [8]. which is our action space. If the model overpredicts the risk level, the cost is 1.0. Otherwise, the cost .s x (y ú) where y is the actual risk level, and is the predicted risk level.
This popular image dataset contains 100 classes, which correspond to our actions, and each class is in one of 20 superclasses. We assign cost follows: 0.0 for predicting the correct class, 0.5 for the wrong class but correct superclass, and 1.0 for a fully incorrect prediction.
CIFAR-100
Results Across tasks, DISTCB achieves lower average cost all episodes (i.e., normalized regret) and over the last 100 episodes (i.e., most updated policies' performance) compared to SquareCB. This indicates the empirical benefit of the distributional approach over the conventional approach based on least square regression, matching the theoretical benefit demonstrated here. Perhaps surprisingly, DISTCB also consistently outperforms FastCB, Both methods obtain first-order bounds with the same dependencies on A and C** which suggests that DISTCB's empirical improvement over FastCB cannot be fully explained by existing theory. The only difference between DISTCB and FastCB is that the former integrates online MLE while the latter directly estimates the mean by minimizing the log loss (binary cross-entroe cross-entropp)..... An even more fine- ane-grained understanding of the benefits of distribution learning may therefore be helpful in explaining this improvement. Appendix K contains all experiment details. Reproducible code is available at https: //gi thub. com/kevinzhou4977/ee /dittss/.b.d//....//////./.
Conclusion
We showed that distributional RL leads small-loss bounds in both online and offline RL, and also proposed a distributional CB algorithm that outperforms the state-of-the-art FastCB, A fruitful direction would be investigate connections of natural policy gradient with MLE distributional- fitting scheme to inspire a practical offline RL algorithm with small loss guarantees, à la Cheng et al. [2022]. Finally, it would be interesting to investigate other loss functions that yield small-loss faster bounds.
This material is based upon work supported by the National Science Founda- tion under Grant Nos. IIS- 8862100 and IIS- IIS-2154711.
Acknowledgements
10

---

References
Alekh Agarwal, Akshay Krishnamurthy, John Langford, Haipeng Luo, et al. Open problem: First- order regret bounds for contextual bandits. In Conference Learning Theory, pages 4-7. PMLR, 2017.
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural complexity and representation learning of low rank mdps. Advances in neural information processing systems, 33:20095 2020.
Alekh Agarwal, Yuda Song, Wen Sun, Kaiwen Wang, Mengdi Wang, and Xuezhou Zhang. Prov- able benefits of representational transfer in reinforcement learning. In The Thirty Sixth Annual Conference on Learning Theory, pages 2114-2187. PMLR, 2023.
Zeyuan Allen-Zhu Sébastien Bubeck, and Yuanzhi Li. Make the minority great again: First-order regret bound for contextual bandits. In International Conference on Machine Learning, pages 186-194. PMLR, 2018.
András Antos, Csaba Szepesvári, and Rémi Munos. Learning near-optimal policies with bellman- residual minimization based fitted policy iteration and a single sample path. Machine Learning. 71:89-129, 2008.
Peter Auer, Nicolo Cesa-Bianchi Yoav Freund, and Robert E Schapire. The nonstochastic multiarmed bandit problem. SIAM journal on computing, 32(1):48-71 2002.
Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for rein- forcement learning. In International Conference on Machine Learning, pages 263-272. PMLR, 2017.
Gabriel Barth-Maron, Matthew W. Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva TB, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributional policy gradients. In International Conference on Learning Representations, 2018. URL https: / /tttps://peenrviiii .eeteeeeeieeee///eeeeeee Syrum?id-SyA1p200b.
Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforcement learning. In International conference on machine learning, pages 449-458. PMLR, 2017.
Marc G Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C Machado, Sub- hodeep Moitra, Sameera S Ponda, and Ziyu Wang. Autonomous navigation of stratospheric balloons using reinforcement learning. Nature, 588(7836):77882, 2020.
Cristian Bodnar, Adrian Li, Karol Hausman, Peter Pastor, and Mrinal Kalakrishnan. Quantile qt-opt for risk-aware vision-based robotic grasping. Robotics: Science and Systems, 2020.
Nicolo press, 2006.
and Gábor Lugosi. Prediction, learning, and games. Cambridge university
Jonathan Chang, Kaiwen Wang, Nathan Kallus, and Wen Sun. Learning bellman complete repre- sentations for offline policy evaluation. In International Conference on Machine Learning, pages 2938-2971. PMLR, 2022.
Jinglin Chen and Nan Jiang. Information-theorede considerations in batch reinforcement learning. In International Conference Machine Learning, pages 1042-1051. PMLR, 2019.
Ching-An Cheng, Tengyang Xie, Nan Jiang, and Alekh Agarwal. Adversarially trained actor critic for offline reinforcement learning. In International Conference on Machine Learning, pages 3852-3878. PMLR, 2022.
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In International conference on machine learning, pages 1096-1105. PMLR, 2018a.
Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018b
11

---

Chris Dann, Yishay Mansour, Mehryar Mohri, Ayush Sekhari, and Karthik Sridharan. Guarantees for epsilon- greedy reinforcement learning with function approximation. In International Conference on Machine Learning, pages 4666-4689. PMLR, 2022.
Christoph Dann, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. On oracle-efficient pac rl with rich observations. Advances in neural information processing systems, 31, 2018.
Simon Du, Sham Kakade, Jason Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in rl. In International Conference on Machine Learning, pages 2826-2836. PMLR, 2021.
Mónika Farsang, Paul Mineiro, and Wangda Zhang. Conditionally risk-averse contextual bandits. arXiv preprini 13577 2022.
Alhussein Fawzi, Matej Balog, Aja Huang. Thomas Hubert, Bernardino Romera-Paredes, Moham- madamin Barekatain, Alexander Novikov, Francisco J R Ruiz, Julian Schrittwieser, Grzegorz Swirszez, et al. Discovering faster matrix multiplication algorithms with reinforcement learning. Nature, 610(7930)47-53, 2022.
Dylan Foster and Alexander Rakhlin. Beyond ucb: Optimal and efficient contextual bandits with regression oracles. In International Conference on Machine Learning, pages 3199-3210. PMLR, 2020.
Dylan J Foster and Akshay Krishnamurthy, Efficient first-order contextual bandits: Prediction, allocation, and triangular discrimination. Advances in Neural Information Processing Systems, 34: 18907-18919, 2021.
Dylan J Foster, Zhiyuan Li, Thodoris Lykouris, Karthik Sridharan, and Eva Tardos. Learning in games: Robustness of fast convergence. Advances in Neural Information Processing Systems, 29, 2016.
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv preprint arXiv:2112.13487, 2021.
Dylan Foster, Akshay Krishnamurthy, David Simchi-Levi, and Yunzong Xu. Offline reinforcement learning: Fundamental barriers for value function approximation. In Conference on Learning Theory, pages 3489-34899 PMLR, 2022.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 770-778, 2016. doi: 10.1109/CVPR.2016.52
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski. Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018
Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Contex- tual decision processes with low bellman rank are pac-learnable. In International Conference on Machine Learning, pages 1704-1713. PMLR, 2017.
Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. Reward-free exploration for reinforcement learning. In International Conference on Machine Learning, pages 4870-4874 PMLR, 2020a.
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory, pages 2137-2143. PMLR, 2020b.
Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of rl problems, and sample-efficient algorithms. Advances neural information processing systems, 34:13406-134 13418. 20
12

---

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline ri? In International Conference on Machine Learning, pages 5084-5096. PMLR, 2021b.
Sham Kakade, Akshay Krishnamurthy, Kendall Lowrey, Motoya Ohnishi, and Wen Sun. Information theoretic regret bounds for online nonlinear control. Advances in Neural Information Processing Systems, 12-15322,,, 2020.
Sham M Kakade. A natural policy gradient. Advances in neural information processing systems, 14, 2001.
Nathan Kallus, Xiaojie Mao, Kaiwen Wang, and Zhengyuan Zhou. Doubly robust distributionally robust off-policy evaluation and learning. In International Conference on Machine Learning, pages 10598-10632. PMLR, 2022.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. Advances in neural information processing systems, 33: 21810-21823, 2020.
J Kolter. The fixed points of off-policy td. Advances in Neural Information Processing Systems, 24, 2011.
Alex Krizhevsky. Learning Multiple Layers of Features from Tiny Images. 2009.
Lucien Le Cam. Asymptotic methods in statistical decision theory. Springer Science & Business Media, 2012.
Chung-Wes Wei Lee, Haipeng Luo, Chen- Yu Wei, and Mengxiao Zhang. Bias no more: high-probability data-dependent regret bounds for adversarial bandits and mdps. Advances neural information processing systems, 33:1552: 15522 -15333, 2020.
Qinghua Liu, Alan Chung, Csaba Szepesvári, and Chi Jin. When is partially observable reinforcement learning not scary? In Conference on Learning Theory, pages 5175-5220, PMLR, 2022.
Thodoris Lykouris, Karthik Sridharan, and Éva Tardos. Small-loss bounds for online learning with partial information. Mathematics of Operations Research, 47(3):2186-22 2022.
Clare Lyle, Mare G Bellemare, and Pablo Samuel Castro. A comparative analysis of expected and distributional reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 4504 4511, 2019.
Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional reinforcement learning, Advances Neural Information Processing Systems, 34: 34:19235-19247, 2021.
Zakaria Mhammedi, Adam Block, Dylan Foster, and Alexander Rakhlin. Efficient model-free exploration in low-rank mdps. arXiv preprint arXiv:2307.03997 2023.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 5 518(7540):529-533 2015.
Anna Montoya, BigJek14, Bull. denisedunleavy, egrad, FleetwoodHack, Imbayoh, PadraicS, Pru_Admin, tpitman, and Will Cukierski. Prudential life insurance assessment, 2015. URL https://kaggh com/ caappetiiionssprudeetttooo -ife- ieffeiinsurance... assessment.
Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka. Parametric return density estimation for reinforcement learning. arXiv preprint arXiv: 1eXiv:1203.3497 2012.
Rémi Munos and Csaba Szepesvári. Finite-time bounds for fitted value iteration. Journal of Machine Learning Research, 9(5), 2008.
Gergely Neu. First-order regret bounds for combinatorial semi-bandits. In Conference on Learning Theory, pages 1360-1375. PMLR, 2015.
13

---

Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline rein- forcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems, 34:11702-11716, 2021.
Mark Rowland, Mare Bellemare, Will Dabney, Rémi Munos, and Yee Whye Teh. An analysis of categorical distributional reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pages 29-37. PMLR, 2018.
Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Mare G Bellemare, and Will Dabney. An analysis of quantile temporal- difference learning. arXiv preprint arXiv:2301.0446) 2023.
Daniel Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic exploration. Advances Neural Information Processing Systems, 26, 2013.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
John Tsitsiklis and Benjamin Van Roy. Analysis of temporal-din diffffrreffeee learning with function approximation. Advances in neural information processing systems, 9, 1996.
Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. In International Conference on Learning Representations, 2022. URL https: //openreview.o net/forum? id=ttr.Jsbteeeee
Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. International Conference on Learning Representations, 2021.
Sara van de Geer. Empirical Processes M-estimation, volume Cambridge university press, 2000.
Joaquin Vanschoren, Jan N. van Rijn, Bernd Bischl, and Luis Torgo. Openml: Networked science in machine learning. SIGKDD Explorations, 15(2):49-60, 2013. doi: 10.1 01144///64119026666 1988.. URL http::ttti.....ee acm. org/ 10. 1145/2641190 2641198.
István Vineze. On the concept and measure of information contained in an observation. In Contribu- fions to Probability, pages 207-214. Elsevier, 1981.
Andrew J Wagenmaker, Yifang Chen, Max Simchowitz, Simon Du, and Kevin Jamieson. First-order regret in reinforcement learning with linear function approximation: A robust estimation approach. In International Conference on Machine Learning, pages 22384-22429, PMLR, 2022.
Kaiwen Wang, Nathan Kallus, and Wen Sun. Near-minimax-oppine risk-sensitive reinforcement learning with evar. International Conference on Machine Learning. 2023.
Ruosong Wang, Dean Foster, and Sham M. Kakade. What are the statistical limits of offline RI with linear function approximation? In International Conference on Learning Representations, 2021a. URL htttt:::/oppnrreieeeeeee
Ruosong Wang, Yifan Wu, Ruslan Salakhutdinov, and Sham Kakade. Instabilities of offline rl with pre-trained neural representation. In International Conference on Machine Learning, pages 10948-10960. PMIR, 2021b.
Yuanhao Wang, Ruosong Wang, and Sham Kakade. An exponential lower bound for linearly realizable mdp with constant suboptimality gap. Advances in Neural Information Processing Systems, 34:9521-9533, 2021c.
Runzhe Wu, Masatoshi Uehara, and Wen Sun. Distributional offline policy evaluation with predictive error guarantees. International Conference of Machine Learning, 2023.
Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-consistent pessimism for offline reinforcement learning. Advances in neural information processing systems, 34:6683-6694, 2021.
Tengyang Xie, Dylan J Foster, Yu Bai, Nan Jiang, and Sham M. Kakade. The role of coverage in online reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023. URL https:://penttt http:///ppeeeeeeeeeeeeeeee ht/ppp/neeeee..te//ttrr//aaaaaaa hett:forrumiie--11jjaaP....
14

---

Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie- Yan Liu. Fully parameterized quantile function for distributional reinforcement learning. Advances neural information processing systems, 32, 2019.
Andrea Zanette and Emma Brunskill. Tighter problem-dependen. regret bounds in reinforcement learning without domain knowledge using value function bounds. In International Conference on Machine Learning, pages 7304-73033 1312 PMLR, 2019.
Andrea Zanette, Martin Wainwright, and Emma Brunskill. Provable benefits of actor-critic methods for offline reinforcement learning. Advances in neural information processing systems, 34:1 13626- 13640, 2021.
Tianjun Zhang, Tongzheng Ren, Mengjiao Yang, Joseph Gonzalez, Dale Schuurmans, and Bo Dai. Making linear mdps practical via contrastive representation learning, In International Conference on Machine Learning, pages 26447-26466 PMLR, 2022.
Tong Zhang. Mathematical Analysis of Machine Learning Algorithms. 2023. http://www. tongzhang- ott.llo-o/tooooooo/oo book html.
Zhengyuan Zhou, Susan Athey, and Stefan Wager. Offline multi-action policy learning: Generalization and optimization. Operations Research, 71(1): 148-183, 2023.
15

---

Appendices
A
Notations
Table 2: List of Notations
S,A,A
State and action spaces, and |A|.. The set of distributions supported by S.
A(S)
d
The expectation of any real-valued distribution i.e. d = Eyndly) {1, ... J for any natural number N.
[
(x, a)
Distribution of Ei=h Ct given In x, an rolling from æ. Q7 (z, a) 27, (z, a) and V E (07, a)]. Optimal policy, i.e., π = arg min, V/T (x1). Without loss of optimality, we take
Q7 a), (x)
n
be Markov deterministic.
Zho Qi. VE TR. Th Th TTi
7 QK V with n², the optimal policy.
The Bellman operators that act on functions.
The distributional Bellman operators that act conditional distributions.
V "M,7", 2T,V*, V (z,0) O
V ( 7" = (x1). V*, Z are defined similarly with x².
The probability of visiting (x, a) at time h.
Coverage coefficient max,, ||dd//drr)) ll
Da(f 9)
Triangular discrimination between f. f.g.
H(f 9)
Hellinger distance between f.g.
lf
KL divergence between f.g.
A.1 Statistical Distances
Let f. be distributions over y. Then,
Da(f g) = Σ (f(y) - g(u))2
f(u) + g(v)
H(J (1 = V
VT(u) - vo(u))
(fI( (1 = Σ (
1 Σi(()) -
Drv (1) 2
The following standard inequalities will be helpful:
H² S Drv S V2H,
2H² D < 4H²
(Lemma A.1)
H DDLL.
Lemma A.1. For any distributions f.g, we have 2H 91 S D& ( & 4H I1 9).
Proof. Recall that
Da (1) -J, f(y)-9(v) VVot+alu
Applying VIG)) < Vi()+o(u)
/2
concludes the proof.
Vi())
16

---

B Modified Algorithms with UAE and for Small Returns Bounds
In this section, we present the O-DISCO algorithm with Uniform Action Exploration (UAE). We also present versions of O-DISCO and P-DISCO for the reward- maximi-na setting (instead of the cost-minimizing setting studied throughout the paper); if SMALLRETURN turned on, we can derive small-return bounds in Appendix J.
Algorithm 4 O-DISCO (with UAE and small return)
1: Input: number of episodes K, distribution function class F, threshold 3, flag UAE, flag SMALLRETURN.
2: Initialize Da,o t- 0 for all h E [H], and set Fo = F.
3: Set op = max SMATL.RETURN else op min.
4: for episode & 1, 2, ..... K do
5: Set f(k) = arg SEF op, (1)
6: Set (xx)) arg op, 7(k) f (x,a).
7:
if UAE then
8:
For each h E [H]. collect Inac di. an, unif(A), che Calenaaaaw)
Pa (2 an,,)) and augment the dataset D4. DA, U (h Cnd else
9:
Roll out * and obtain a trajectory #1.k, 01,k1 01.55 T H. HHHK: CH..v For each [H], augment the dataset DA, Dh, U (T Ch,ki Ch In+11,))
11:
end
12:
For all () [H] x F, sample v ( a') and a' arg Jo+ (1 a). where (In,i, ah,i, Ch.s: x),;) is the i-th datapoint of Dr.k- Also, set 2h, Ch,s + Vn and define the confidence set,
Fx = -EF
h Σ 1=1
log fn fn This an,)) 2 max JEF
log Fr Ih,s ap,i) - 73, 4h Iml
13: end for
14: Output: uni
Algorithm P-DISCO (with small return)
1: Input: datasets D DH, distribution function class F, threshold 3, policy class II, flag SMALLRETURN,
2: For all (h, fT E [H] x F x II, sample w fn-1 Th+1 where
2fa (Kh,i, On,ss Ch,br ,//,)) is the i-th datapoint of DDp Then, set 2h, = Ch + w and define the confidence set,
Fx
Σ log fn (1 Whis ap,i) 2 max
Σ
an.s an,)) - 78, Vh !n] log Fn (
JEF
Set op max if SMALLRETURN else op min.
4: For each π E II, define the pessimistic estimate fr = arg op (annma))(i((xx
5: Output: मे = arg оPrEπ Eaws(z)) [f [R),
17

---

C Proofs for DISTCB
Lemma C.I (Azuma). Let (Xi) be sequence random variables supported on [0, 1), adapted to filtration (FilseNy For any E (0, 1), we have w.p. at least 6,
N
Σ
B[x Fi-1 - Σ X, V log(2/6),
(Standard Azuma)
N Σ 1=1
N
B[X
Fi-1l 2
Σ 1=1
X, + 2 (og(1/6).
(Multiplicative Azuma)
Proof. For standard Azuma, see Zhang [2023, Theorem 13.4]. For multiplicative Azuma, apply [Zhang, 2023, Theorem 13.5] with À = 1. The claim follows, since T-exp(-X] 2
Theorem 4.1. For any 8 E (0, 1), w.p. at least 6, running DISTCB with 7 10A
40A(C* +log(1/)) has regret scaling with C = > 112 Regretion (K))oo((11)))
k= minlaeA C(I),a).
Regre () < 2 AC* Regreb (K) log(1/5) 2300 (Regrel (K) + log(1/5))
Proof of Theorem 4.1. First, recall the per-step inequality ReIGW Foster and Krishnamurthy [2021, Theorem 4], which states: for any and y 2 2A, we set p = ReIGW,(f.1) then, for all I e [0, 1]4, we have
Pla(a)(((a)) - f(a))) S $4 (())((a)) + 77 T. p(a) (fa))
f(a){((a)
where a* arg min, f(a). For any & E [K], applying this to f ((6) p= = Pk and O(ss,), -), we have
5A Ölsk, ax) +77 f(k) ax) + C(ak, ak)
(F2) ax) - ((kk,a)))
Eau (()) C(sk, (sk)) S
Fas
5.4
Exg
Ö(sk, ax) +7:Da ax) || C' ax))
1= 1
(Eq. (41))
Since DA < 4H2, we have
K
Eas
(a ax) ll ((kk,))))
Ea. H (H((C((s, ak) (()) at))]
4
A=1
8
H22 (C(sk, ax) g(0) (sx, ( 8lo((1/8) (Multiplicative Azuma, since 122 E [0, 1))
< 8 Regret (K) 10log(1/0))
(Foster al. [2021, Lemma A.14J)
Hence, we have
K
5.4 Σ Ea. [O(s,, ax) 70y (RegRegtt (K) + log(1/8)).
Emc [[(s,, ax) -
<
Finally, recalling that 1/(1 -)) < 1 28 when e < 2. and the fact that 54 < 2, we have
K
K
10.4
Ea (C(SE, - C(sk, TT" ((
Eag [C(ss, x (sx))] + 1407 (Regretiog (K) + +log(1 -log(1)))..
7
k=
18

---

By Azuma's 's inequality, we have
(sky as) C(sk, at" (sx))
Rap [C(sk, ax) C(sk, x (s:))] log(1
K
20A4
Bay (8k; n (sa)) + 1400 (400gReegee (K) + log(1/6)) + 2 log(116)
7
k=1
40A
1
(C + (log(1/8)) + 140r [Regreting (K) + log(1/6)) + 2 log(1/8).
7
(Multiplicative Azuma)
Now set
140 Regret (K)ggg((1/))) 10A
40A(C* (10g(1/6))
Case
is
when
40A(C'
10A, ie. (C* + log(1/0))
140 (K)
280 (Regreg (K) + log(1/8)). we have the above is most
4(C" + log(1/6)) + 1120.A (Regretigg (K) + log(1/6)) + log(11/)) < 2240 (RegreRg (K) + 1og(1/0) + 2 (1g(6))).
Case 2 is when the left term dominates, then the bound is,
2 4480A(C* [Rg((//)) (Regret (K) + log(1/6)) + + (og((1/))
2/ 13440AC* Regrel.nn (K) loglll//////// 44800 log (1/6) + 2 2log(1/1)
232 AC* Regret (K) log(1 (1gg(111)) /6) + 134v/Alogg /8) + 2 log(1/8).
Putting these two cases together, we have the result.
19

---

D Distributional Bellman Completeness in low-rank MDPs
The goal of this section is to show that, under mild conditions in low-rank MDPs, there always exists a function class with bounded bracketing number that satisfies the distributional BC condition. First, let us recall the low-rank MDP In this section, we show that linear MDPs automatically satisfy the distributional Bellman completeness assumption.
Definition 5.6 (Low-rank MDP). transition model PA : X x - ((X) has rank d there exist unknown features 0% : X A + Rd, 14 X' + I such that Pa x a) = o% (2, a) 1 (x') for all x, a, x'. Also, assume maxz,a ( (x, a)ll2 12 I and [I gdyt llz I |[9l|so v for all functions g: X' R. The MDP is called low-rank if Ph. low-rank for all h E [H].
Suppose that we have function class & such that of E Φ for all h, i.e. a realizable function class. For example, in linear MDPs, this is automatically satisfied since know Ø' a priori, so + is the singleton with ø² Having a realizable & class is standard for solving low-rank MDPs [Uchara al., 2021, Agarwal al., 2023].
In what follows, let a = =axa,,...... Z7 (2 x, a) denote the maximum density/mass value of the loss-to-go distributions. Note that 2 1 always since the mass H + 1 is deterministically placed at zero. If we further know that Z is discretely distributed, then a 1. If Z is continuously distributed, assume is bounded.
We consider the function class in Eq. (2), which we reproduce here:
Flll f(2 2. (0(z, a), a), w(z))
E b. [0. Rd
(3)
s.t. max (=)22 avd and max (ō(x, w(2) saj 2,0,5
The next lemma (Lemma D. shows that this function class satisfies distributional BC. Lemma D.1. Flin satisfies distributional BC (Assumption 5.1).
Proof. We denote ||/||>- = maxz,r. f( xr a). For any fa+1 E Flin h++1- we have | S by the construction of h+-1 Then, let T be either the distributional Bellman operator distributional optimality operator, the following equalities hold for the appropriate a' a'(x') based on T
/ Pr(x'l 12, / Pr(c æ, an+iil - aa',a((z')) dæ de fx h
TD fn+1 x,
()), / (a((")) / Pr(c x. cl x', (x'))) dedr'
Since J& Pra (e a) (, c cl d'(')))()))))))))))))))( de S 11.n+tllos 11o. we know that
[|wn(2]]le & II V avd.
We further note that
max (x, a), wx(z)) = -max m fnt1 x2 S llfntallo aa 2,0,2
Also note that o7 E by realizability. Therefore, TD fnt1 E Flin, which is the distributional BC condition.
D.1 Bounding the bracketing number via discretized rewards
We now bound the bracketing number of Jlin under a discretization assumption that costs and costs-to-gos can only lake M many discrete values as discretizing the reward space, and can be shown that this discretization error small for regret or PAC bounds [Wang et al., 2023, Section 6]. Structural assumptions are necessary to bound the
evenly spaced grid. This can be interpreted
20

---

complexity of Fht and such discretization assumptions are common in practice, e.g. C51 [Bellemare al., 2017] and Rainbow [Hessel et al., 2018] both M 51 which works well in Atari games. After discretizing, we can consider w as a mapping from [M]. the discrete set on M elements, rather than from the interval [0, 1]. Note also that since 7 are discrete, we have a =.
Now, let s > O be arbitrary and fixed. Recall that the foo bracketing number is equivalent (up to universal constants) the &oo covering number, so we will work with the latter. Let B(r) denote the d-dimensional ball of radius r (in 62). Recall that the s-covering number (in (2) of functions [M] + Br) scales as ((rr/e)d)) ddM) Let We be such the smallest cover. We can build Eo cover of Fl as follows: C. = {(z, 2) + (Ø(x, a), w(2)), w We, )
To check this и cover, consider any f Fh.. corresponds to some 6 and w. Let be the neighbor of in W. and let 2)) (d(z,a), w"()) indeed f' E C.. Then, for any æ, 0, 2, we have (ő(x,, a), w(2) w ==(=))) I l,((,, ()lall()))) w'(=)1l2 S e. Hence, C₂ is an lo cover of size (( ((VV/s))M [(]), and so we have shown that log (e, Thi, Iloo) s
Odd log(d/E) + log (
Linear MDPs: Recall that in linear MDPs, know the true o- and so | 1. Thus, the bracketing number simply OdM log(d/E)) in linear MDPs.
Summary and comparison with regular BC: In summary, under the assumption that rewards are discretized, we know that low-rank MDPs automatically have distributional function classes that satisfy distributional BC and have bounded bracketing numbers. Furthermore, recall that [Wu et al., 2023] showed that Linear Quadratic Regulators (LQRs), with deterministic transitions, also have function classes that satisfy distributional BC and have bounded bracketing numbers. Thus, distributional BC holds for the most interesting cases covered by the standard Bellman completeness, e.g., linear MDPs, low-rank MDPs and LQRs. Since learning conditional distributions is statistically harder than learning the conditional mean, we need to pay the price in assuming reward/transitions satisfy regularity assumptions to bound the bracketing number appropriately.
21

---

E
Generalization Bounds for Maximum Likelihood Estimation
This section reviews generalization bounds for the maximum likelihood estimator (MLE). We adopt the same sequential condition probability estimation setup as in Agarwal et al. [2020, Appendix E], which we now recall for completeness. Let X be the context/feature space and y be the label space, and we are given a dataset D = {( 36 3i) from martingale process: for i = 1, 2, nJ, sample D; (1 M1:i-1 and y: ~ p(- æz)) Let (a, v) = ply xx) and we are given realizable, i.e., f E F. function class F X' y - A (IR) distributions. The MLE is an estimate for f that maximizes the log-likee lieeekkeihoodd objective over our dataset:
FALE = arg max
Σ log f(2,, yc).
JEF
For guarantees to hold for general hypotheses classes F, use the bracketing number to quantify the statistical complexity of F [van de Geer, 2000].
Definition E.1 (Bracketing Number). Let & be set of functions mapping X' - RR Given two functions 1, u such that ((x) S u(()) for all x E X, the bracket [t, u] is the set of functions g E g such that [(x: I (xx) I u(2) for all x X. We call u] e-bracket if ||u I|] S Then, the -bbccketing number of G with respect to |1-11. denoted by Ni 9, II-II) is the minimum number of c-brackets needed cover G.
Since the triangular discrimination equivalent to squared Hellinger up to universal constants, we now prove MLE generalization bounds in terms of squared Hellinger.
Lemma E.2. Let fi X + A()) and f2 X x y + E. satisfying supzex Sy f2(1, y)dy s s. then for any distribution D E ((Y), we have
ENDD (f(x)) I| f2(2, ))] s 10-11 1) 1
-2g E. Da exp log log(fi(z.))//
Proof. This follows from the proof of Wu et al. [2023, Lemma C.I].
Lemma E.3. Fix E (0, 1). Then w.p. at least 6, for any F, we have
BaND, (H(((x,)) 11 (x,
6nd Dl +
> log (z., + (N, (llm)))))
(4)
am1
Rearranging, we also have
Σ log (f(r) S 3ne + 4log (ogg(N(e,,F,ll) (ll
1lx)
(5)
Proof. We take an e-bracketing of F, {14, us] i 1, 1, and denote J {uy = 1, 1.2,...) Applying Lemma 24 of Agarwal et al. [2020] to function class JF and using Chernoff method, w.p. at least 1 5, for all F. we have
log E exp(L(f(D). D') D'
(llg))))2 llc))
(6)
(ii)
(1)
Now, fix any f E F and pick f E F as the upper bracket, i.e., f S f. Now set L(f, D) E 11/2log(f (x ff)/f(xr, vi)). Then the right hand side of (6) is
log(S (x,, uz) fx, vi)) + 2log (Nog((N (c,FF (1,=)))))))))))) (1x)
log(f* (2..)=(())) + 2 2 (c,FF, (l,)))))))))))))))))))))))))))))
22

---

On the other hand, since H is a metric, we have
E H (F(x,)),
(I(I)))).( + H (Mo.)))))
(H
DrDD
IDD
H² ((..))
2
E
+2
IDDi
(i1)
(iv)
For (iii), by the definition, we have f(x, v) f(r, v) E [0,c] for all x,
(iii Σ E
H²
1 [(x, f(x, w) w) - s 2ne ((x,,)) < Σ E z~D;
2
D-D;
For (iv), we apply Lemma E.2 with fi = I and ʃ2 f (chus 8= = 1 + and get
(F(x.w))
(iv) =nc -
log
B
exp 2 log
2 (2,)
f=1
=nd - Σ log
E E.y~D.
exp
2 log
(f(x.W)/((.v))
n Σ
(Ma.w)/f(..)) n
=nd 2log
E
exp
log
D.
i=1
= + (i).
By plugging (iii) and (iv) back we get
"
EE E (f(a.)) S 6ne + 4 (i). DND
Σ
Notice that (i) & (ii), so we complete the proof by plugging (ii) into the above.
We first state the MLE generalization result for finite F.
Theorem E.4. Suppose FF is finite. Fix any 8 E (0, 1), set 3 = log((F118) and define
f x= =={EF:: Log f(r.)))) max feF
f(zn y.) 48
Then w.p. at least 8, the following holds:
(1) The true distribution is in the version space, ie., 1 F.
(2) Any function in the version space is close to the ground truth data-generating distribution, for all J E F
Σ BLLDD (1F( < 223.
= |F]. For (1), apply Eq. (5) to f fMLE that f* F. For (2), apply Eq. (4) and note that the sum term at most 43. Thus, the right hand side of Eq. (4) is at most (6 + 8 + 8)3 223.
Proof. These two claims follow from Lemma E.3 with < = 0, and so (c, -1111111))))))
We now state the result for infinite F using bracketing entropy.
23

---

Theorem E.5, Fix any 8 (), ser B = log Iog .F,1 (lco)/5) and define
8= デ= =={EE= Σ log flz u.) 2
max feF 1=1
flan v() vo) -73
Then w.p. at least 1 8, the following holds:
(1) The true distribution is in the version space, i.e, f E F.
(2) Any function in the version space is close to the ground truth getta-genrratig distribution, for all f E
s 283.
Proof. These two claims follow from Lemma E.3 with e = 1/w/y/ For (I), apply Eq. (5) to fo f to see that f E F. For (2), apply Eq. (4) and note that the sum term is at most 73, Thus, the right hand side of Eq. (5) is at most (6 + 14 + 8) = 283.
Confidence set construction with general function class
In this section, extend the confidence set construction of O-DISCO and P-DISCO to general F, which can be infinite. Our procedure constructs the confidence set by performing the thresholding scheme e-net of F. While constructing e-net for F is admittedly computationally hard procedure, this is still information theoretically possible and our focus in O-DISCO and P-DISCO show that distributional RL information-theoretiiaae leads small-loss bounds.
We first define some notations. Let Fi and FF denote a lower and upper E-bracketing of F, Le., for any f E F, there exists an e-bracket [ such that for all h, f% S fn S Sh with I E Ft, F,,s S Ft. Recall that a lower bracket E F may not be a valid distribution, but since elements of F map non-negative values,
assume has non-negative entires well. Also, so for e small enough, g is normalizable. Hence, define g(2 xr a) = a [ar, a) g(( x, a) as the normalized version, which is valid distribution that we sample from.
we have a. (2, a) = f g)(2 x, a) 2
Now, consider any martingale { tn,i, Ch, Sesseeeeeeeee hejar which could be the online data up to episode i or the offline data (consisting of i...d. samples). We define the MLE with re- spect to lower bracket element as follows. For any h [M].9 E F4, R E II, sample _..
g.z yh, 9N+1(Ih,,s T (2%,,)) and 2h,i h as,
Ch.i + sha define the MLE solution for (g, x) at time
MLE R
arg max JEF
log fü (2 9.
Th.i, ah,i)
1=1
Also, define the version space with respect to the above MLE as.
F gg = {SET: Σ log (= Thei an.))
log MLE (=1. (= Jh,s an.i) -
We now prove key result that implies that Th. 1 falls into the confidence set F,t,
Theorem F.I For any 8 E (0, 1) and suppose n 2 2. Then, w.p. at least - ō, for any E [H) E F.f+ E Ft, E II. we have
Σ log gn (ofttt"" Inhs an.s) log Th" Shez (2 This, an.s) S log(e" NN /
where The"" cha + whi and whi Phanl (h+14a1))))
24

---

Proof of Theorem F.I. Consider a E-bracketing of F where S 1/n S 1/2; we will study each element and conclude with union bound. For any lower bracket 1 and upper bracket in the bracketing (note 1, need not correspond to the same bracket), Recall that 1 a) == f 4 (= I, a), we have 1 S a since I is lower e-bracket of distributions. Therefore, we have
un (=1, Tn an.))
)
24 (21,5 1: 2 Thii ab..))
E
exp
log
II Evn,
T Dhis, on.s)
1T Th ans))
2m
where vh,i is the distribution of data from i-th round and time h. Note that en, a C, c,,'') ( C)Ch(e x, (. x; a) for some distribution dn, a). Now focus on each i, so for all i, we have
un(chi 2h,iaan.
Eva.s
Th Th Ihis an,)
[ Phole,,,c,d Di L', x(x)) Jow un,(c, of x, (.a)ttlll x1,x(x)))
unle + y a)
J. Jans dn.lee un(z 2,0)
[ phile, o' |2,0)tnt2(2 - -cl'', Sew Dnole, x' (2 2),a)n+2(2 xx'x)x²)))
0,
dns(o, un(2 2, (1)
1,a,2
1
1+
23
1
+
1-8 5 =1+-
Therefore,
un (= Tn,i, ans)
]
E
exp
log
(1 4/n)" se.
1,t Thit (Zh,e Thii an.i)
Thus, by Markov's inequality, w.p. at least
we have
un(= L:t
Thii, ah
log
< In(e"/00 /6).
Th 11+ (2/, Thai an.a)
To conclude, apply union bound to get this result for all brackets.
For the remainder of this section, assume the policy class 1I is finite. However, it is possible to extend our results using policy covers in the Hamming distance; in that case, log II| would be replaced by the log covering number entropy integral of II [as in Zhou al., 2023, Kallus et al., 2022]. We note that for the online case, we rely on the assumption that for any f E F we have T E II, where recall that T = arg min, fr (I, a). This is because T is not a contraction so we cannot operate with T.. directly and instead operate with _ We highlight that this assumption is automatically satisfied in tabular MDPs, since the whole policy space is finite, and log [T| = O(X log(A)) lower order compared log of the bracketing entropy of Ftab. which O(x2 A2). In contrast, in non-distributional methods such as GOLF, the regular Bellman optimality operator is contraction so standard Lipschitz arguments for covering go through. We note that is also possible to construct covers of F in the Hellinger distance, but the metric entropy of Frab seems to be on the same order as its bracketing entropy.
We now describe the version space construction for general F, first for the online setting, Fix any k. and define the set
(EF: Σ log fn (sfte" Thii, an.))
log M
ft. Ihii an.i) 3
Then, construct the version space
F = fo E Vh [H]]
25

---

Theorem F2. Fix any ò E (0,1) and suppose Assumption 5.1. Set / = (o Ni(KK) FF.11
ll /11)). Then, w.p. at least 1 - 8, the following holds:
(1) The optimal cost distribution is in the version space, i.e,, Z E Fs.
(2) For all f E Fk and h E [H)
Σ Ai (fr
II Th (n+1(2n) S <603.
Proof. First, we want to verify that Z* E Fk. Let f be the lower bracket of 7* and set g MLE ft.5" E F; note x² T². By Theorem F.1, we have Σ log MLE " (sh |
Th,s, (h.i) - Tog D
Tha) S ( Therefore, noting that Z D sh+ shows that 7 E FF bs for every h, implying that T Fk. For the second claim, fix any E F& and [H] Then,
Th 7 2 TT'
Ei (F an) 11 Th (1+i (x)))))) an))]
B. (H)((()) ll Th D (n+il an))
Σ E HT' an) ll (on, an)) + (TT D Phii
11 Tr ".' Ihtt an))
2(28,3 + 3kz).
The 3 comes from Theorem E.5, and for e, we used the fact that H² S H S TV, and
E
(V D "Tht (un, an) Tr' (1ni
1=1
Th" T T Ph an)an) Th' T fntz(2 (m,
E
1=1
4:
Eyg's Eulc,a (oh, an) Phozl cl x', x' (x')) fht1(2 f -1aa'
35 3kE,
since for any I, a, we have J. Ti (=1 x, - fn+1(2 x, a) S 3e. There are two cases.
Th (z x, x. 2 fn+1(2 2,,), then Pht1(z x, - fh+1 (z x) S (1- -) -
a, - fn+1(z a. I 22fn41(z a, since (1- < 1+ 2e. If Phi1 (z x, a) <
fnt1(z x, a, then frt1(2 a) 7l fz (z 2, s fht1(z xx - fh x, S e. Thus,
Jg max(22 x, a), =) S S 25 f 2, = 3z. Thus, setting = = 1/K gives
k
F [ an) || T (nti ax)) < 593
For the offline setting, fix any and define its general version space as,
=x = (f F: fn Fy Vh [A]]-
26

---

Theorem F.3. Fix any 8 E (0,1) and suppose Assumption 5.1. Set B = log( H [1]
N(((n(D)))1 FI,1l 118)) Then, w.p. least 6, the following holds for all policies II:
(1) The policy cost distribution is in the version space, i.e., Z" Fx.
(2) Any function in the version space has bounded triangular discrimination with the ground truth data- generatinn distribution, i.e., for all f E Fx and E [H]
Eyyn (I,(,, an) Th Sh fut1(an, aw))
6003N11
Proof. The proof is the same as in Theorem E.2, but instead of x/, we fix any
G The lp distributional eluder dimension
Let S denote any input space (for example, we will later instantiate S = X' or S = X x A). Let denote set of functions mapping from S R. Let be of distributions on S.
Recall the definition of e-independent sequence (of distributions) from Jin et al. [2021a].
Definition G.1 (fz-independent sequence). A distribution v E D is (e, -inddeeennde of sequence
{d()) "("",...,ddee dd))) C D if there exists v E v such that [,, v > and also v (E +) S
Note that the definition is sequences of distributions, which generalizes the original definition sequences of points from Russo and Van Roy [2013].
We now generalize the above definition for the general Cp norm. Definition G.2 (fp-independent sequence). A distribution v E 'D is (e, fy)-inndpenddent of a sequence
{d()) d(m) } C D if there exists vi E V such that [E, > and also > |B v/ < sP.
Using the definition independent sequences established so far, we define the é, distributional eluder dimension.
Definition G.3 (lg- distibbbbillii eluder dimension). For any p, define the e distrribbuti eluder dimension (denoted by DE,, D)) the length of the longest sequence (d ...., d(d)) D such that there exists e' 2 E, such that for all E [d], (+)) (E', (dd, independdnnd of (11) d(t
Of particular interest to us is the é case. We show that the & eluder dimension is dominated by the le eluder dimension Jin et al. [2021a].
Lemma 5.4. For any V, D and O, we have DE, DE (V, D,&) S (V,
Proof. Since v E, < Σ (zg), we have that any witness (long independent sequence) for E₁ is also a witness for E2. So, the maximum length of the bg witnesses is longer than the & witnesses. Liu al. [2022. Proposition 19] obtains an analogous result for the non-distributional eluder dimension of Russo and Van Roy [2013].
We now prove the key pigeonhole result for the & distributional cluder dimension.
Theorem 5.3. Let C == SupdED, [Faf] be the envelope. Fix any K N and sequences
f()) f(K) C V, d(1),, d(")....,d d(k) C D. Let B be a constant such that for all k E [K], we have, k-1 Eas f [0]] < 3. Then, for all E [K], we have
Σ
Eag /00 S inf (e, + Blog(C/e)) + ke).
0<e51
Proof. For any I C D, and 0 <<<<< let ((,,,e)) denote the number of disjoint subsets of T such that each subset is e-dependent of v, i.e,, for all such disjoint subsets of it not the case that v is (e, E1) inddeeeendd of each subset.
27

---

Fact 1: For any E, if [(0)) > for some k E [K], then L(d(M) d(ick-)), < 8/E.
By definition of == L(d(h), d(lik=1),)) E), there exist disjoint subsequences G(),...,, ..... ((L) of
d( such that each subsequence ((() satisfies Exer(o) [af(() > Therefore, summing all subsequences, we have Ls < premise. This proves Fact 1.
E ii=1 [ |(h|| S B, where the 3 inequality comes from the
Fact 2: For any and any sequence (u(")...,,NN L(vl),, (l.f)1)))))))) e) 2
D, there exists j E [sc] such that
= [(m 1)/ DE, (V, D)
J = 0, the claim is vacuously true. Otherwise, consider the following algorithm for finding the j:
Step 1) Initialize ((1) [ G()) [b()] and let j = J +
Step 2) If ((ا)) is e-dependent on all of ((), E [J], then the claim is proven and terminate.
Step 3) Otherwise, there exists some ((,, i E [J] such that vt) is e-independent of it. Append wG) to ((),, ie., (+))) 5()) [u(t]] Increment j j + 1 and go back to Step 2.
Hence, need argue this process terminates Step 2 before j gets to + 1. We prove this by contradiction: assume j gets
1. Let i [J] be such that (()) has the most elements (break ties we have that [B(1 Z k/J - DE (V, D,E) > DE, (V, D,E), where we've also used the definition of J. By construction, (6)) is an E-eluder sequence, Le., it is sequence such that each element e-independent of its predecessors. However, this is a contradiction because its size is greater than DE₁ ( (V) D, e). Therefore, this process terminates Step 2 for some j. which the witness for proving Fact
arbitrarily). Since Ei= @ S (
Fact 3: For any and E [K], we have Ei= Ea [ > s] s (3c-1+ 1))111))BBF DE, (V. +1.
Fix any E and A E [K]. Let [d(i)) .........edd d(in) be all the elements of d(1:k) such that Egco ((1 > for t i1...., Ax- By Fact 2, there exists j [x] such that L(d(')) d(busss)) E) 2 (U,
By Fact 1, we have L(d(t)) d( < B/c. Finally notice that L(d(t)), d(m)-1), e) < L(dis), d(I:))),) e) since adding more elements can only create more E-dependent-of-v disjoint subsets. Thus, combining these incqualities, we have [(x 1)/ DE, (V)) B/c. This implies x S (3e-1 + 1) DE, (V, D, + 1, which proves Fact
[(x
Finishing the proof
Fix any k E [K] and w > 0. We have
k
k
Σ
E Re f( p(0)
Σ
E p(t/|
yydy
=
1=1
k
> >g]dy Earos f10/
< kw
k
kww
Σ
E Exco f(0/
>y >y]dy
< kw +
(18/v+ + DE, (V. D, + 1]dy
(Fact 3)
< kw [(8/y + 1)DE, (V. D.w) + 1)dy
(Monotonicity of DE,)
kw + + 11C + lo dBlog(C/w).
(d = DE (L, w)
This completes the proof.
G.1 Bounding V-type & eluder dimension in low-rank MDPs
Theorem G.4 (Bound of l2 distributional eluder for low-rank MDPs), Suppose the MDP low-rank MDP. Let V c. X + [0, 1] be any class of functions mapping X [0, 1). Suppose
28

---

D= ( (x) π E II) for some h E [H]. Then, we have
DE2 Dee) O(dlog(d/e))
Proof. If / = 1, then D is a singleton. Hence, DE2 (v, D) S 1. Hence, suppose h 2 2; set h = - 1 and we will focus on 4xt1 in the remainder. Suppose { {()) ) is any sequence
such for all k E [T]. we have that (d(k), (()) is (e, -indeeennnd of its predecessors. For any k; set Ex --1 (o (Z an gx0 (3h, an)] + AL. Then, we have
B (()) (In+1) = Eg0 f
0% (I an) " d'dp (xn+1)) (21+1)
gus (T [
(
Jn11
Be (Fn, an) (n)lsssll| / (
Focusing on the second term,
Bat [rW) ((azz))) Ad
1h41
Thus, we have shown that
A Σ
Ege p(k) (In+1) S ( an) an)l
Baa [(( (xn++))) + Ad.
Then, by the independent sequence assumption, have
T
T
Tec
Bgas f (IN+1)
E o% (In, anll
EEg [+++2))))) Ad
k=1
k=1
T
| di (2h. an) an>ls (e vtd)
E (E (In+)))) 5)
k=1
2s
[|acs o7 (wns (n)la
(1 = s2/d)
k=1
T
< T
|E 200) of (In, Al
25 / T log 7/da)
(elliptical potential)
T doog T/+2).
=² /d)
For a reference of the elliptical potential, see Uehara et al. [2021, Lemmas 19&20). Rearranging, we have vT < 2 /og + T/2), which implies
T T 4d log(1 T/-2).
By applying Lemma G.5, we have & log(1 + 4d/2). This concludes the proof.
Lemma G.5. Let c1, C2 > 1 be constants. Let 0 be solution to I c1 log(1 o2z). Then, necessarily have x & 6c, log + c102).
Proof. Using change of variables B a. we have the inequality is equivalent B S log(1 + B ci c102). Take exp of both sides to get cxp(B) S aB + where a C2 From Step of the proof of Russo and Van Roy [2013, Propasition 6]. we have B < e- (log(1 + a) + log(e/(e - ())) <
3(log(1 + c102) + 1). Hence, c1 33(log(1 c102) 1).
29

---

G.2 Bounding Q-type & eluder dimension in tabular MDPs
Theorem G.6 (Bound of &g distributional eluder for tabular MDPs). Suppose the MDP tabular MDP. Let v C A + [o, 1| be any class of functions mapping X x A to [0. 1|. Suppose 'D be any set of distributions. Then, we have
DE₂ vp D) O(SAlog(SA/e)).
(8)
Proof. Suppose {d((), (()) kE(T) is any sequence such for all & E [7], we have that (d( f(k) is (e, fa)-independdnt of its predecessors. Since the MDP is tabular, A-ddmeensional vectors. For any set Eか = =xE (d()) + AZ. Then, we have Eg (1) a) ( ((A)) I ||d") || |1sx
interpret d(8) ((h) as
Focusing on the second term, we have
12 Si
(d())) T(t))) + ASA.
Thus, we have
T
k-
Tec
<
Ege (1) (I, I
Σ
(ld
BBa (2,, ())) + ASA
k=1
h=1
[d ||
E + VASA
< 25
|||(
( = (2/SAA
k=1
2
|d(k)) 1ي__.
S 2: TV SA log 1 + +T/E2).
(elliptical potential)
Rearranging, we have vT 2 SAlog(1 T which implies T < A + T/s²). Then by applying Lemma G.5, we have T s. 24 (1g1111 +4S4/2)) This concludes the proof.
30

---

H
Proofs for Online RL
H.1 Preliminary Lemmas
Lemma H.1. For any policy n, conditional distribution d and E [ we have
Th d(z,a) -h T))
Th d(z,))= = (000)TT700 a))
Proof.
Th:Dd(r, d(x,a) a)
E VTT (aa)ll
-endduttaaarrrhhtesdddooamae ay 57
= Cn(2,, a) + =
Cn((za) + -B [d(x',
Th d(r,a) a) Ey y ddx,,))
" ming d(Ø' a) al) 1).c 1 s/)
= C)(I,a) ( a'arg min, d(as ( a) ,3/13/)
= =Ch(I,a) (a.a), d(o' a)) [d(a', ,d)]
Ch(I,a) EEr mind( a')
= Ty d()
Lemma H.2 (Performance Difference Lemma (PDL)), For any f :x x A A- RR" and policies T, nt we have
H Σ
V
E. E (2)) ( (f)f((()))))
Ey
Th' f1+1 an) fn (2hh
Proof. We proceed by inducting on the following claim: for all h = H + H, H......
H
V (mn) Sn (2b. x²)
B
TH' Seti (t) S Se(2,,
t=b
The base case of H + is trivially true as everything is 0. Now fix any and suppose the IH at + true. Then
Ví (#x) - ('))
= E.n,n (ch + Vh+1(#n41) (#1+1) - f141 + I n') - falan x')
ETT,n [W+ (xN+11) - fh+1 n')] + Brzzlen + Jht1 (1 fn (x)nn
By the IH, the first term is equal H Et=h+ Er,,, (T Seti 0r) Se m²)] The second term exactly En,,, TT sh+1 (1 an) - Jn (2 (') which concludes the proof.
31

---

H.2 Proof of Small-Loss Regret and PAC Bounds
Recall that we defined the function class and distribution class, for each h, as
(1I) ((z,a) (2, II) V (<x,a) - (a( a) || T (()
(10)
Also, define the 'V-type' analogs as follows, which will be useful for PAC instead of regret bounds.
DA, (II) f (z) II}
(11)
A = {z {++ Barufii(( ()) a) T ((.
Let us also overload notation for the eluder dimensions
(e) max DE nv, (T)T))
(E), max DE, ( h,.. Dh,e (II).
Before we prove the following main theorem, couple of remarks are in order:
1. Recall that by Theorem G.6, we have (e) S AoSAAoo(SS log(SA/e)) and by Theorem G.4, we have (E) I Od log(d/s)). This shows that the Eluder dimension in terms in Theorem 5.5 are appropriately bounded.
2. In Appendix D, we showed that distributional BC (Assumption 5.1) is satisfied in low-rank MDPs and the log bracketing number is bounded by O(dM log(d/e) + log Φl) where & is realizable class for o² This shows that the BC assumption Theorem 5.5 is satisfied and appropriately bounded for low-rank MDPs.
Taken together, these two points imply that we have small-loss PAC bound for low-rank MDPs: concretely, we have Va V* I dH /Av-
AV log (b
d d log [+]
K
We now prove the our main result for online RL: Theorem 5.5. We will prove the result with general function classes, so we will replace the |F] by its loo bracketing number, i.e. = log ( (o,))).
Theorem 5.5. Suppose DisuBC holds (Assumption 5.1). For any E (0. 1), wxp. at least 6, guarantees the following regret bound,
rumning O-DISCO with B
Regreto.Disco(( DISCO (K) & 160 / KV (E,(1/K) K) log( K) B- + 18000 ( K) log( K)3.
UAA TRUE (Algorithm 4), then the learned mixture policy म is guaranteed satisfy,
AV* (1/K((KK) log(K)8
18000//²ADE; (1/K (og(K)
V - VV < 160H
K
K
Proof. For shorthand, let (n, = (x) || The D Sh+1 (xr, a)) and A z(k)
DH 15, an)). Notice that since (x)) ((0) (z, a), (n, (f( ((() (z, Th J (z, a)) well.
arg min, JJt+ (z, (z,a), a), we have
Th D (k)
Jht1 (z) =Th DD.
D (K)
By Theorem F22, we have the following two facts for all & [K].
(i) Optimism: min, ((() (x11 a) < V* (since Z* F2) and
(ii) Low training error: for all ht we have
If UAE=FALSE. Eick Es (5, an) < 240/3
UAE=TRUE. Eicr Egy [E (6 an)]] < 2403.
The 240 comes from the constants of Theorem F.2 and the fact that Da b) S 4H² ((, b) for all distributions 6.
32

---

Now, fix any episode & E [K].
Va V
$ V
min fi (Z)
(Fact (i)
Exg
Th Jhtl (2n, an) - -h (on, ot (an,
z(k)
(2n))
(PDL Lemma H.2)
Eggh
Th T.tth c(N)
J+t1 (2p, an) - Fh (Zh an)
(Lemmaa H.1)
h=1
Ex [a (sn, an)
dn, an)
V B (dn, (coh, an)
(Eq. (A2))
H
H
4eV** + 17H
Σ t=h
Bo VBar ( an)
(Lemma H.3 and E, (K), an) S V*)
S Vee+** 17HH VHA.
(*)
S V4HHV=* A + 5HAR
< HV + 2 HT/H + 5H
In * we used Cauchy Schwartz. Setting 1 = 4vH and rearranging, we have
V=' 2V 16// A + 10// S 2V* + 26H
Plugging this into *, and noting 104e + 17 < 300, we have
V= V v 8eV + 300 A HAR
Thus, summing the instantaneous regrets over all episodes, we get
V
8eV' 300H/ HAR
SeKV + 300H Σ A
H 2 A
55 HKV* Σ A + 18H
Σ
A
Last step: bounding Σg Ag. In this final step, we invoke the pigeonhole property of the eluder dimension, as proven Theorem 5.3. Note that the precondition Theorem 5.3 is satisfied by Fact (ii) mentioned at the beginning of this proof. Also, since the triangular discrimination is always bounded by 1, we have that C' in Theorem 5.3 at most 1, and we will also pick = 1/K.
On one hand, if UAE=FALSE. then,
K
11
K
Σ A
Σ
Enp I an)] < 1000/ DE, (1/K)Blog(K)
1=1
h=1A=1
On the other hand, if UAE=TRUE, then, we use the V-type analogs,
K
H
K
A
En (Zn, on)
1=1
1=1
k=1
K
Eyp (aaaai On (xt, u)]
1=1 1=1
< 1000AH DE (1/ / log(K)
33

---

This concludes the proof for both the regret and PAC bounds.
Lemma H.3 (Self-bounding lemma). Let I E F and let a be any policy. Let us denote on (2, a)
(a(fa(z,a) |I Th fn+1(2,))) a)). Then, for all h E [Lf], for all Th, Ohe have
H
fa(In, S eQi as) + 4H
Σ
(or at)]
t
Proof. We prove the following refined subclaim inductively: for all h E [H], for all Th, ah we have
t-h
[a an S
Σ
Eg (a), ae) + (xt,
(IH)
H
t=h
For H + 1 this is trivially true. Now fix any and suppose IH is true for h + 1. By Eq. (A22), for any
h, The O4. have,
fa an) Th. (14,( an) S
4TH fn+1 an) (a an) òn (2 an)
S 4TH Ju+ (2oh, on (on an) 5% (#h, an)
1 Th Intilan, an) ) +HH 1) ap).
(AM-GM)
H
In particular, we have that
fa(an, an)
+ H
Th FN+11 ax) 22 (2 on)
1 H
(Ca(Tnaan) an) + Bes- [T (IN+1, + 2H (pn,
t-h-1
1
in (xh, an) + Pg
Σ
1
B (7 (2 ax) + d ac))
-=
(IH)
2H (IA, ax);
which proves the inductive claim. Noting that (1+1H/HH)) < e, we have proven the lemma.
H.3
Regret Bounds for Tabular MDPs
Theorem H.4 (Small-loss regret for tabular MDP). Suppose the MDP is tabular with X states and assume Assumption 5.1. Fix any E 10 and set B = log H K F1/8). Then, w.p. at least 1 8,
Regret (K) E O(H/XXKV' H2XA3..
In terms of H, K, K scaling, our bound matches that of GOLF [Xie et al., 2023] and only a H factor looser than that of the minimax lower bound Õ(v X AK). The key benefit over prior bounds is that our leading term scales with the minimum
of the problem V*. For example, if V Rv 0, O-DISCO attains Olog K) regret while uniform regret bounds are lower bounded by 2 ((VK). Compared to the minimax-optimal UCBVI [Azar et al., 2017], one weakness of our theorem that it needs a F satisfying BC. Fortunately. in tabular MDPs where cost only revealed the last step from a known distribution, we can choose Fiab as described in Wu et al. [2023, Lemma 15] to automatically satisfy BC. By extending our theory via bracketing entropy (Appendix F), we can derive that Feab yields 3 = (Xx A22 log( (og(XAl AAHK) We note that if costs are unknown but discrete, possible construct a BC function class with / scaling as (xx2, A2 log(n AHK/5)) A where is the maximum number of possible cumulative costs.
34

---

Extension to linear MDPs The Q-type dimension captures Linear MDPs when squared loss is used by exploiting the fact that the bellman residual is linear in o" a) [Jin et al.. 2021a]. However, since our function class is the set of triangular discriminations, rather than the Bellman residual, we find that the Q-type dimension does not immediately capture Linear MDPs unless regularity assumptions are made. For instance, we believe that Linear MDPs are captured by the Q-type dimension if we assume that (z æ, a) is lower bounded, i.e., the value distribution is sufficiently smooth.
I Proofs for Offline RL
Theorem 6.1 (Small-Loss PAC bound for P-DISCO). Assume Assumption 3.1. For any å E (0, 1), w.p. at least 1 1-6 running P-DISCO with 3 = log [[]] /ô) learns a policy that enjoys the following PAC bound with respect to any comparator policy म
11:
C#V
C0
V V 9H
N
N
Proof of Theorem 6.1. For shorthand, let 8% (z. Dz a) Th
(1+1(2,)) and A
E Ex (2 a.)]. Also, let f(x, π) Bt) ((2,
By Theorem F.3, we have the following
facts, for all ı E II,
(i) Pessimism: V/ < fr (2). π)) (since Z" E Fx) for all π E IT, and
(ii) Ey, 18% (In, ag B'NNI for all h where Theorem F.3 and the fact that DA I 4H² certifies
that 3' = 2403 is sufficient.
With these two facts, we can bound the suboptimality of क as follows:
V V
$ (21, 7) - V
(Fact (i)
$ (x V
(Policy selection scheme Algorithm (Line 4))
"
Σ
E7
Ih (2h, 7) - The Jhr (In, an)
(PDL Lemma H.2)
Ex 44f (I an) + 6% (I an)] V B (I an)]
(Eq. (Ag))
H
4eV + 17H
Σ 1=
Ea (Ie: ar)] E [ an]]
(Lemma H.3)
4=1
< V 4 + 1 HA*
< AVHVE + 5
Finally, we can bound A by change of measure,
H
E
1h. an)
Byg (8n (In,
1=1
T H -'N1
(Fact (ii)
Therefore,
C#V z
C
V - v
4H
N
N
35

---

Extension: Small-Return Bounds
In this section, we show that O-DISCO and P-DISCO can also be used obtain small-return bounds. Compared to the algorithms presented in the main text for minimizing cost, we simply have replace min with max (and vice versa) for maximizing reward, i.e., see Appendix and enable the SMALLRETURN flag. The proofs are also largely the same, with slight changes to the first few steps.
Theorem J.1. Assume Assumption 5.1 and suppose we want to maximize returns (instead of minimize cost), so enable the SMALLRETURN flag. Fix any 5 E (0,1) and set B log H K F Fi/)) and 3' 603. Then, w.p. at least 6, running O-DISCO (Algorithm 4) with UAE FALSE yields the following small-loss regret bound,
Regreto.Discol (K) & 5H LSECKK)) + LSEC instead UAE TRUE, the outputted policy ग enjoys the following small-loss PAC bound,
(12)
AV (K)) 1811 ALSEC,(K))
- v $ 45HV
K
Proof. Adopt the same notation in the proof of Theorem 5.5. By Theorem F.2, we have the following two facts for all [K],
(i) Optimism: V max, f((h) a) (since Z E Fi) and (ii) Eich Egg (6n, az)] )] S 3r for all h. UAB=TRUE, then On is sampled from unif(A) rather than n², i.e., we have Etek B,, (nn=r(a))mmi((()) certifies that 3' = 60,3 is sufficient.
( (n, an)l S B", where 8 S 3. Theorem F.2
Fix any episode [K]. Then,
V y**
< max (( (II, - V**
(Fact (i))
Σ F
sh (In, ( (Fn, Th (Th,
p(())
(PDL Lemma H.2)
H
Egth
((()) 5h
(In, (xn,
(N) -T= fiil (2n ( an)
(Lemma H.1)
E
( dn, V Byn (8n. (2o,, an)
(Eq. (A2))
" 4eV** 17H Σ t=
By 18, ar)] / (6n an)
(Lemma H.3 and E, (5 an) $ V*)
< 4eV** +17HA, HA,
< + H
Thus, summing the instantaneous regrets over all episodes, we get
K
V'
4eV 17H HAr
1=1
4eKV 17H Σ A H AB
(Cauchy-Schwand
k
55 HKV* Σ A
+
18H Σ
AK.
The bounds for AS are the same as in Theorem 5.5.
36

---

In some sense, the proof for the small-returns bound is actually easier than the small-loss bound. Recall that in the cost-minimizing setting, we needed to perform a crucial Cauchy-Schwand step to rearrange terms at the step labelled + However, in the reward-maximizing setting, we simply bound V=" V*, without needing to rearrange terms.
Theorem J.2. Assume Assumption 5.1 and suppose we want to maximize returns (instead of minimize cost), so enable the SMALLRETURN flag. Fix any 8 E (0, () and set B |o [ Then, w.p. at least 5, P-DISCO (Algorithm 4) learns policy स such that for any comparator policy
II, we have
C N
30H²CC33
V V & 99H
N
Proof of Theorem J.2. Adopt the same notation as the proof Theorem 6.1. By Theorem F.3, we have the following two facts, for all π E II,
(i) Pessimism: (x T) < V" (since Z E Fx) for all x E II, and
(ii) Ev, 16% (In, as) )] S 3' N for all h where 3 608.
With these two facts, we can bound the suboptimality of म follows:
V V
f1 ()),
(Fact (i)
v
(x1,
(Policy selection rule in Line 5)
Σ Ex The Jh+t (2h: an) - (2 #)
(PDL Lemma H.2)
H
Ex [4Sh (Kh, an) + + (oh, an)]
Ea [o (nh, an)]
(Eq. (A2))
h=1
From here, the same argument in the proof of Theorem 6.1 finishes the proof.
K
Experiment Details
Experiment Settings
In experiments, outlined in Foster and Krishnamurthy [2021], learning rate at each time step is set to "/- = "Nut" where "t and are hyperparameters. We use batch sizes 32 samples per episode, and the King County and Prudential experiments for 5, 000 episodes while the CIFAR- 100 experiment runs for 15, 000.
For each dataset, we select the hyperparameter configuration with the best performance for each algorithm. As we report two metrics, performance over the last 100 episodes and over all episodes, we choose the best hyperparameters for each metric well. While it is often the same hyperparameters that give the best last 100 episodes and all episodes results for a model, that is not always the case. We the WandB (Weights and Biases) library to sweeps over hyperparameters.
Oracles
For our regression oracles, we use ResNet18 [He et al., 2016), with a modified output layer (so that the output is suited for 100 prediction classes) for CIFAR-100. and simple 2 hidden-layer neural network for the Prudential Life Insurance and King County Housing datasets. For DISTCB, the oracle's 's output layer has size ACC where is the number of actions and C the number of potential costs. This is reshaped so that for each action, there are predictions associated with each potential cost, which then have a softmax function applied to them represent cost probabilities. For SquareCB and FastCB, the output size is A because there is just single prediction associated with each action. As per Foster and Krishnamurthy [2021], I1, a sigmoid function is applied this output layer. All experiments were implemented using Py PyTorch.
Datasets
We now provide an overview table as well as additional details and context to our setups for each dataset. Note that the number of items in each dataset in the table is the count after preprocessing.
37

---

Datasets Items
Dataset
Number of Actions
Number of
Costs
CIFAR- 100
50, 000
100
Prudential Life Insurance
59, 381
King County Housing
20, 148
100
101
Table 3: Overview of the three datasets and their experimental setups
Prudential Life Insurance This dataset is from the Prudential Life Insurance Kaggle competition [Montoya et al., 2015]. It featured in Farsang et al. [2022], which inspires our experimental setup. The risk level in [8] directly determines the price charged to the customer. Thus, we can consider the chosen risk level as the action taken. If the model overpredicts the risk level, we get a cost of 1.0 because this is considered over charging the customer and getting sale. Otherwise, the model's prediction is charging too little for the customer. To reiterate, the cost in this case is .1 * (y D) where y the actual risk level, and is the predicted risk level.
King County Housing The King County housing dataset also used Farsang et al. [2022]. An interesting part of the setup is that the cost construction in the case of not overpredicting differs from the Prudential experiment, even though they' re both effectively about predicting a price point. Here, the model's chosen price is considered the gain, which is why the cost is 1.0 minus the chosen price. On the other hand, in the Prudential experiment, the cost is a linear function of the difference between the chosen value and the actual value.
CIFAR-100 For the CIFAR- 100 experiment, we use the training dataset of 50, 000 images as our dataset. The inclusion of the superclass is critical, it lets delincate 3 possible costs that DISTCB can learn. Without the super class, the cost construction would be pure binary correct vs. incorrect. If this were the case, the ability to test the effectiveness of learning the distribution would be nullified. The distribution would just be whether an action is correct or not, which means our algorithm would essentially be predicting the mean directly.
Results
The largest advantages DISTCB had over the next best algorithm were in the Prudential experiment, with DiSTCB having .086 advantage over the last 100 episodes and .045 advantage over all episodes. While the gaps were not as large for the other two datasets, they are still statistically significant and further showcase the benefit of distribution learning.
38


Reviews:
[{"review_id": "gxGJiWqMID", "reviewer": "Reviewer_9rQf", "summary": "The paper's main concern is to theoretically understand why distributional RL achieves good performance. They consider three different settings: contextual bandit, online RL with an optimistic algorithm, and offline RL with a pessimistic algorithm.\n\nIn each case, they use MLE to learn a distribution over the unknown cost (bandit case) and loss-to-go (RL case). The key technique that enables the proof is the relating a new notion of distributional regret to the regular regret; this is done by manipulation of distributional divergences. It appears that the key across settings is that distributional divergence gives more fine-grained control over value/cost differences compared to only looking at means. \n\nThis enables them to provide small-loss bounds in each case. They validate their findings empirically for the contextual-bandit case.", "questions": "1. The triangular discrimination technique seems central to converting distributional divergence into bounds on value difference. Intuitively, what are the main factors that enable the small-loss term to show up when distributional approaches are used? Are there ways to make this term appear without distributional RL? If not, is there reasoning for why it is not possible?\n\n2. How restrictive are the realizability and Bellman completeness assumptions made in the analyses? Do you have a sense of how the techniques could extend to violated assumptions?\n\n3. Would it be easy to generalize the results to linear MDPs? What obstacles may arise?", "limitations": "Yes", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Significance and Originality:\n\n1. The problem considered is important and relevant to practitioners. The authors do a good job in the introduction of motivating the theoretical conundrum.\n2. The tools developed by the authors around distributional divergence are novel to the best of my knowledge.\n\nQuality\n\n1. The results appear correct. And the settings considered are comprehensive across contextual bandit, online and online RL.\n2. The empirical results display improved performance over well-chosen benchmarks.", "weaknesses": "Clarity:\n\n1. The regret decomposition using triangular discrimination, while novel, is intricate and difficult to interpret intuitively. If further intuition were provided, this would help inspire future theoretical work.\n\n2. The writing is notation-heavy and takes a while to parse through and keep track of notation.\n\n3. Relating to point 1, it would be nice if there were further connections to the practical applications around risk-sensitive RL, which are the primary motivating examples for distributional RL. This would help bridge the gap to practitioners.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive and constructive comments! Please find our responses below.\n\n**Triangular discrimination:**\nIntuitively, triangular discrimination bounds give finer control over estimation error than traditional $L_2$ bounds from non-distributional methods. For example, for some target conditional distribution $g(y\\mid x)$, imagine we want to estimate its conditional mean $\\bar g(\\cdot\\mid x)$. On one hand, square loss regression would learn a function $\\hat f_{sq}:\\mathcal{X}\\to\\mathbb{R}$ such that $|\\mathbb{E}\\_x[\\hat f\\_{sq}(x) - \\bar g(\\cdot\\mid x)]|\\leq \\\\\\|\\hat f\\_{sq}(x) - \\bar g(\\cdot\\mid x)\\\\\\|\\_{L\\_2} = \\mathcal{O}(1/\\sqrt{N})$. On the other hand, suppose we learn a conditional distribution $\\hat f\\_{dist}:\\mathcal{X}\\to\\Delta(\\mathbb{R})$ with MLE. Using triangular discrimination, we can obtain a *self-bounding* inequality $|\\mathbb{E}\\_x[\\bar{\\hat{f}}\\_{dist}(\\cdot\\mid x) - \\bar g(\\cdot\\mid x)]|\\leq \\sqrt{(\\mathbb{E}\\_x\\bar g(\\cdot\\mid x) + D\\_\\triangle(\\hat f\\_{dist}(\\cdot\\mid x)||g(\\cdot\\mid x))) \\cdot \\mathbb{E}\\_xD\\_\\triangle(\\hat f\\_{dist}(\\cdot\\mid x)||g(\\cdot\\mid x)) }$. This is in fact the *implicit inequality* that can be derived from Eq. $\\Delta_1$ on Page 5. By standard MLE generalization results, we expect $\\mathbb{E}\\_x D\\_\\triangle(\\hat f\\_{dist}(\\cdot\\mid x)||g(\\cdot\\mid x))=\\mathcal{O}(1/N)$. Thus, if $\\mathbb{E}_x\\bar g(\\cdot\\mid x)\\approx 0$, we expect the bound to converge as $\\mathcal{O}(1/N)$, which is faster than squared loss's $\\mathcal{O}(1/\\sqrt{N})$ rate. This separation between squared loss regression and MLE is actually fundamental, and there already exists a lower bound in the CB setting, see Theorem 2 of Foster et al., 2021 \"Efficient first-order contextual bandits...\" To summarize, we can only obtain this key self-bounding inequality with MLE, and not squared loss, and this is the key intuition for how we obtain small-loss bounds. \n\n**Realizability and Bellman completeness:**\nSince BC is stronger than realizability, we focus our discussion on BC, which we posted in the global response. In short, we can prove distributional BC for tabular, linear, low-rank, and LQR MDPs. We will add this to the paper to address your important point.\n\n**Generalization to linear MDPs:**\nAs remarked in the global response, DistBC indeed covers linear MDPs. Moreover, we proved another new result inspired by Reviewer 8JhK's comments, which shows the LSEC (an analysis tool used in Appendix G.2) is bounded by the Bellman eluder dimension. Since low-rank MDPs have Bellman eluder dimension $\\widetilde{\\mathcal{O}}(d)$, these two new results imply that our small-loss bounds also hold for low-rank MDPs (and thus also linear MDPs), further generalizing our results!\n\n**Risk-sensitive RL:**\nRisk-sensitive RL is indeed well-motivated for distributional RL, so there is no conundrum there. The long-standing conundrum about distributional RL is regarding the risk-neutral setup: when optimizing expected returns, why can learning the distribution then computing its mean perform better than learning the mean directly? (By Bellman equations, all we need are the expected returns, so why should we learn the distribution?) We answer these questions with small-loss bounds, which *converge faster* than bounds from non-distributional methods; to the best of our knowledge, our work shows the theoretical benefits of distributional RL for the first time. With that said, we believe the techniques developed in our paper could also be useful for deriving small-loss bounds in risk-sensitive settings, and leave that as future work. \n"}, {"review_id": "L9U55Y2Ruu", "reviewer": "Reviewer_L483", "summary": "This paper explores the benefits of distributional reinforcement learning (RL) and provides a mathematical basis for its advantages. Traditional RL approaches focus on learning the mean loss-to-go, but recent developments have shown that learning the entire loss distribution can lead to improved performance in various tasks. However, the theoretical understanding of why and when distributional RL works well has been limited. The paper introduces the concept of small-loss bounds, which are instance-dependent bounds based on the minimum achievable cost in the problem. By optimizing over distributional confidence sets constructed through distributional Bellman equations, the proposed algorithms achieve small-loss regret bounds in tabular Markov decision processes (MDPs) and small-loss PAC bounds in latent variable models. The paper also presents a distributional contextual bandit algorithm and an offline RL algorithm with a novel robustness property. Empirical results demonstrate the effectiveness of the distributional RL algorithms in challenging benchmark tasks. ", "questions": "None.", "limitations": "The paper includes a brief discussion of the limitations.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "This paper investigates an important problem in theoretical RL and deepens our understanding of its benefits based on rigorous mathematical analysis. It proposes new distributional algorithms for contextual bandits, online and offline RL settings and provides corresponding small-loss bounds. The work presented in this paper is novel and original, to the best of my knowledge.\n\nThe paper is well-structured, building from the simple setting of contextual bandits and then extending the analysis to more complex RL setups. The paper clearly states the assumptions, theorems and the proof sketches in each section and provides formal algorithms wherever required. ", "weaknesses": "The readability of the paper is poor owing to the large amount of text/math. It seems the authors have modified the spacing between lines and headings in some parts of the paper in order to adhere to the page limit. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your positive comments! Distributional RL is indeed quite notation heavy, and since we show its benefits in all three settings of CBs, online RL, and offline RL, we necessarily need to use notations from all three settings. We hope that Appendix A's table of notations can serve as a convenient index for searching notations, and we'll be sure to add more text descriptions to improve readability. "}, {"review_id": "ejRKGRs85s", "reviewer": "Reviewer_nAEh", "summary": "This paper establishes the regret bounds of the proposed distributional methods from the perspective of the small-loss bound, where the regret upper bound depends not only on the number of iterations but also the optimal value. Based on this, they can show that regret could be lower under those environments whose optimal policy corresponds to a higher value function.\n\nThis paper provides three new distribution RL approaches for three different RL settings. The first one combines ReIGW and MLE to solve contextual bandits. The second one and the third one use the similar concept, which apply confidence intervals on the estimation of likelihood, to solve online RL and offline RL, respectively. They show that the proposed method can achieve competitive regret bound theoretically. The numerical results also demonstrate the empirical performance of the distributional approach for contextual bandits.", "questions": "I agree with the authors that there is very limited theoretical understanding of distributional RL, and therefore overall I can appreciate the contribution of this paper which offers a new viewpoint for understanding distributional methods in RL. That said, the algorithms proposed and analyzed in this paper are all of MLE style, and somehow they are quite distant from the mainstream distributional RL methods (e.g., C51, QR-DQN, IQN, etc). While the results in this paper are nice to have, I am not sure by how much these analytical insights could benefit the understanding of these popular distributional RL methods. In other words, one of my concerns is that the usefulness of small-loss bounds is tied to the specific algorithms presented in this paper. It would be very helpful if the authors could comment on the connection between the small-loss bounds and other common distributional methods.\n\nSome detailed questions:\n- Line 55: The paper claims that triangular discrimination is a novel approach for decomposing the regret, but this idea has been adapted from [Foster and Krishnamurthy, 2021.\n- Line 131: The cost distribution should be some distribution in [0, H-h]? Similar issue occurs for the distribution of the loss-to-go in Line 146.\n- Lines 127&148: The notation \\bar{C} and \\bar{Z} is not defined.\n- Line 152: How to define the sum of two distributions (or two random variables)? Do we assume independence here?\n- In Algorithm, line 4 and Algorithm 3, line 3: What is $\\bar{f}_1$ in line 3? \n- I am not sure how to compute the inner max of the confidence set. It seems that if the policy is not given, the \\mathcal{F}, which is a function of {\\pi}, is hard to compute.\n- There are many hyperlinks pointing to the wrong positions. \n- Line 626, the notation in this inequality is misleading, does f_1 only depend on x or (x,y)?\n\nI am starting with 5 and would be willing to raise the score if the authors could address these questions and the issues mentioned above.", "limitations": "The authors do not describe any specific limitation of this paper.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The concept of small-loss bound is very interesting and appears novel in the context of distributional RL. This new perspective indeed offers a promising way to theoretically understand why distribution methods could empirically achieve better performance in the existing distributional RL literature.\n- This paper is thorough in that it considers three important RL settings: (i) contextual bandits, (ii) online RL, and (iii) offline RL. This paper rigorously shows that the distributional MLE methods have provable benefits in terms of regret for a wide variety of RL problems.\n- The idea of combining MLE and confidence set is interesting and different from the traditional idea of considering the confidence interval of mean return.", "weaknesses": "- Page 7 seems to be missing. This makes the upper half of Page 8 somewhat difficult to read.\n- The regret bounds depend on some factors about the size of the distribution class. For example, in contextual bandits, the regret bound depends on Regret_{log}(K). Similarly, in online and offline RL, the regret bounds depend on $|\\mathcal{F}|$. The finiteness of distribution classes appears to be a fairly strong assumption in RL.\n- Accordingly, another concern is on the assumption that there exists good distribution classes that satisfy Bellman completeness. While it is mentioned in Section 5 that this assumption could hold under some special tabular MDPs, it is unclear how far this argument could go in the more general RL settings. \n- The numerical simulations only discuss the result of contextual bandits. While typically I would not complain about the simulations in a theory paper, I do think that experiments on both online RL and offline RL could be very helpful in understanding the connection between small loss bounds and the empirical regrets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks so much for your constructive feedback. Please find our responses below. \n\n### For the Weaknesses section:\n1. We used the allowed 1 page PDF to upload the contents of Pg. 7, which got accidentally cut. We sincerely apologize for the inconvenience. Luckily no crucial material was accidentally omitted.\n2. **a)** Depending on $\\text{Regret}_{\\log}(K)$ for CB regret is fairly standard, e.g. see Theorem 1 in Foster and Rakhlin, 2020 “Beyond UCB: …” and Foster and Krishnamurthy, 2021 ”Efficient first-order…”. Essentially Theorem 4.1 translates the decision-making regret to the regret of online learning of the distribution, which can be bounded for particular classes and learning algorithms. The crucial point is the dependence on $C^\\star$. \\\n**b)** We consider infinite distribution classes in Appendix F. We mention this at the bottom of page 6 but can certainly advertise this extension much better in the main text. We focused on finite classes in the main text to keep it brief and since finite classes are commonly considered in RL theory. Nonethless, we did extend to infinite classes, as we agree it is much more realistic. The complexity measure we use for infinite classes is the bracketing entropy as it is the standard complexity measure for MLE, cf. Van de Geer, 2000 \"Empirical Processes in M-estimation\". For example, if $\\mathcal{F}$ is a linear function class with features of dimension $d$, then its bracketing entropy is $\\mathcal{O}(d)$. \n3. Please see our discussion on BC in the global response. In short, we can prove distributional BC for tabular, linear, low-rank, and LQR MDPs. We will add this to the paper to address your important point. Thanks for urging us in this direction.\n4. As you write, our primary contributions are theoretical. Running additional experiments for the RL case may be difficult as our DistRL algorithms are based on GOLF and Bellman-consistent pessimism (BCP), which are version space methods that are NP-hard to run. With that said, we'd like to highlight two possible directions for practical versions of the algorithm. First, our confidence set construction is for deep exploration; if the problem only needs shallow exploration, we can adopt a cheaper exploration strategy such as $\\epsilon$-greedy, cf. Dann et al., 2022 \"Guarantees for $\\epsilon$-greedy…\" Second, a follow-up to BCP successfully implemented its main algorithmic idea and showed state-of-the-art results in offline RL benchmarks (Cheng et al., 2022 \"Adversarially trained…\"). Since our offline RL alg shares similarity with BCP, we believe our work provides strong support for empirically investigating a distributional version of Cheng et al., 2022. We leave the implementation and benchmarking as promising future work. \n\n\n### For Questions section:\n**Mainstream DistRL algorithms:**\n1. C51 is actually very similar to MLE, modulo a projection step that is needed due to discreting values. If $Z_{\\tilde{\\theta}}$ is the learned distribution from the last step, C51's update aims to learn $Z_{\\theta}$ that minimizes $KL(\\Phi \\mathcal{T} Z_{\\tilde{\\theta}}|| Z_\\theta)$, where $\\Phi\\mathcal{T}$ is the projected distributional Bellman operator. Since minimizing KL is equivalent to MLE, C51 is essentially doing MLE with projection, so our insights may well apply to C51-style methods.\n2. Quantile-regression (QR) methods such as QR-DQN & IQN minimize the pinball loss rather than maximizing log-likelihood. While we use guarantees in the squared Hellinger from MLE, QR gives guarantees in the Wasserstein distance. It is interesting future work to explore the theoretical benefits of QR for decision making.\n\n\n**Detailed Questions:**\n* Line 55: Foster and Krishnamurthy, 2021 study solely CBs and their analyses do not immediately lead to RL bounds. In contrast, our novel techniques (e.g. self-bounding Lemma G.4) enable us to go beyond CB and prove to the first small-loss bounds in offline and online RL (in non-tabular settings). (Our CB results are not significant in view of Foster and Krishnamurthy and are only meant as expositional warm-up for our RL results, which _are_ novel and significant.)\n* Line 131: The costs-to-go are in [0,1] rather than [0,H-h] since we work under the normalized cumulative costs setup, i.e., costs and cumulative costs are normalized in [0,1] as in Jiang and Agarwal, 2018 “Open Problem: The Dependence…”. This setup allows for sparse costs and is a more general than assuming costs to be normalized in [0,1] up to rescaling by H, i.e. in our setup, 100% of the total cost can be obtained at a single step, while in the traditional setup, only 1/H-fraction of the total cost can be obtained each step.\n* Line 127 & 148: The \\bar notation, which we defined in Line 128, denotes averaging over a distribution. We will remind the reader in a few spots.\n* Line 152: Yes; by + we meant convolution here, that is, the distribution of the sum of independent draws from each distribution. We will simply avoid this notation and describe the distribution explicitly. Note this is the standard distributional Bellman operator, cf. Definition 4.8 from Bellemare et al., 2023 \"Distributional Reinforcement Learning\". \n* Our online RL confidence set does not depend on policies (it is defined wrt Bellman optimality operator). Our offline RL confidence set *does* depend on policies (it is defined wrt policy's Bellman operator). We discuss their computational complexity in the uploaded PDF (Pg 7). \n* We will ensure hyperlinks are fixed for the final version. (Splitting main text and supplement broke all links.)\n* Line 626: Thanks for catching. Where we write \"$f_1(x,y)$\" on the rhs we meant the density of the distribution $f_1(x)$ evaluated at $y$. We will add the assumption that the density exists and give it a notation.\n\nThanks again for your constructive feedback. We will incorporate all these comments in the camera-ready, and please let us know if you have any additional questions. \n"}, {"review_id": "shWpVvtF35", "reviewer": "Reviewer_8JhK", "summary": "This paper explores the benefits of distributional RL and provides a theoretical understanding of its advantages through the lens of small-loss bounds. The authors propose distributional algorithms for contextual bandits, online RL, and ofﬂine RL, and prove that these algorithms achieve small-loss regret bounds. \n", "questions": "see weakness.\n\n\n", "limitations": "yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "I appreciate the attempts to understand the practically powerful algorithm and also the provided empirical simulations, which should be the promising direction in the RL theory study. \n\nThe obtained bound scales with the minimum loss so can outperform standard results when it is small. This result seems to be the first line of work in the literature on RL. \n\nThe paper is also very well-written and easy to follow.", "weaknesses": "1 The algorithm of RL setting still follows 1) reduce the out-of-sample target (immediate regret) to the in-sample loss by complexity measure (SEC); 2) maintain a version space by effective estimation of in-sample loss and choose the optimistic estimator to handle the difference between V^* and V_f (min_a f_1 in this case). The technical novelty is rather limited since similar treatments have been developed in the literature. Also, as far as I know, there are some proofs based on Banach fixed point theorem for the distributional RL, which typically assumes that we can always solve the sub-optimization problems. These results should have supported the effectiveness of the distributional RL to some degree.\n\n2 The proposed algorithm relies on the (distributional) Bellman completeness assumption, which is hard to satisfy because it is non-monotone (adding a new function can violate this assumption). While I understand it is common in the literature to handle the double-sampling issue, it limits the contribution of this work if the goal is to provide a theoretical understanding of the DRL algorithms. I am a little bit confused indeed because, in my mind, the MLE analysis does not require a completeness assumption. Does the distributional setting lead to unique technical challenges here so you have to make such an assumption? In particular, I am wondering whether it is possible to replace it with the trajectory average technique (at a cost of worse regret bound) as in [2,3]. \n \n3 The definition of linear SEC. I am wondering whether there are more non-trivial examples captured by it. For linear structure, my understanding is that the square comes from Cauchy-Schwarz to relate $|\\phi(z_t)^T \\theta_f|$ to $||\\phi(z_t)||_{\\Sigma_t^{-1}}$\n\nand $ ||\\theta_f||_{\\Sigma_t} $. The square also exists in the definition of the standard eluder dimension. One exception is the $\\ell_1$ eluder dimension developed in [1] (see the end of section 5) where both the out-of-sample prediction error and the in-sample training error are linear. I am wondering whether there is any direct application or corollary between them. Another minor comment is that you may also mention the decoupling/eluder coefficient in [4, 5] where such a reduction treatment is first proposed and generalized by them and you can easily bound the complexity measures via SEC.\n\n**The generality of the framework is indeed my main concern**. Actually, I personally have obtained a very similar first-order result with OMLE. I think we have a similar technical issue that we need to have a result as in line 778. My choice is to define the eluder dimension as in [7] (so, we can call it linear eluder dimension just like linear SEC) but I could not find more interesting instances covered by it. \n\n4 An interesting point is that you choose to decouple the original target with two types of complexity measures. In contextual bandit, the immediate regret is decoupled to another out-of-sample target plus some decoupling cost (that is why we need an online oracle). It turns out that this type of complexity measure can be sub-optimal (especially in model-free cases, see [6] for examples). I am wondering whether this choice is only for a convention or for technical requirements because otherwise, we can definitely apply the eluder techniques to reduce the immediate regret to the in-sample error as in the RL case (that is why we only need an offline oracle that looks back at the k-1 samples collected so far). \n\n\n[1] When Is Partially Observable Reinforcement Learning Not Scary?\n\n[2] Contextual decision processes with low bellman rank are pac-learnable\n\n[3] Bilinear classes: A structural framework for provable generalization in rl\n\n[4] A Provably Efficient Model-Free Posterior Sampling Method for Episodic Reinforcement Learning\n\n[5] GEC: A Unified Framework for Interactive Decision Making in MDP, POMDP, and Beyond\n\n[6] A note on model-free reinforcement learning with the decision-estimation coefficient\n\n[7] Eluder-based Regret for Stochastic Contextual MDP\n\n\n\n\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks so much for your constructive feedbacks and please find our responses below.\n\n1. **a)** You are absolutely right that our RL algorithms use the global optimism & pessimism ideas from GOLF (Jin et al., 2021a) & Bellman-consistent pessimism (Xie et al., 2021), and that is actually exactly our point: by simply changing squared loss regression to distributional RL via MLE (just like how C51 simply changed DQN's square loss regression to MLE) we can obtain much faster finite-sample rates in both online and offline RL. Notably, by only changing the loss from square loss to distributional loss, our theoretical _ablation_ proves that distributional RL is indeed responsible for the faster small-loss bounds. Additionally, we highlight that most of our RL results are the first small-loss bounds in these settings (e.g. offline RL and online non-tabular RL), which we believe to be a significant novel contribution in its own right. \\\n**b)** Regarding prior distributional RL (DRL) results with Banach fixed point, these convergence results are *asymptotic* to the best of our knowledge, and hence do not show any *benefits* of DRL compared to vanilla non-distributional RL, i.e. DRL's bounds are not any better than those from vanilla RL. Our key novelty is that our small-loss bounds have a *faster convergence rate* ($\\widetilde{\\mathcal{O}}(1/N)$ in small-loss regime) than the typical $\\widetilde{\\Omega}(1/\\sqrt{N})$ rates attainable with non-distributional methods; in this sense, our work is the first to illustrate concrete theoretical *benefits* of DRL compared to vanilla non-distributional approaches. \n2. **a)** You are correct that MLE by itself does not require BC (see eg. Theorem E.4), but since we perform MLE recursively in a TD-fashion (i.e. the target of the MLE depends on the learned distribution from the last step) we need BC to guarantee all the MLEs will succeed (see eg. Theorem F.2). In other words, the reason BC appears here is exactly why it also appears in the non-distributional setting: squared loss regression itself does not require BC, but TD-style methods like GOLF and FQI do require BC to guarantee squared loss regression succeeds each time. \\\n**b)** Please also see our discussion on BC in the global response. We can prove distributional BC for tabular, linear, low-rank, and LQR MDPs. We will add this to the paper to address your important point.\n3. Thank you so much for pointing us in the direction of the $\\ell\\_1$-eluder dim [1], which inspired a new result we can now prove. The LSEC is indeed bounded by the $\\ell_1$ *distributional* eluder dimension in a manner similar to how SEC is bounded by $\\ell_2$ distributional eluder dimension, cf. Proposition 7 of the SEC paper. We can also show that $\\ell_1$ distributional eluder is bounded by $\\ell_2$ distributional eluder, cf. Proposition 19 of [1]. Altogether, we've shown the following chain: $\\text{SEC}(\\{f^2:f\\in\\Psi\\},\\mathcal{D},K)\\leq \\text{LSEC}(\\Psi,\\mathcal{D},K)\\leq \\text{dim}\\_{\\ell_1}(\\Psi,\\mathcal{D},1/K)\\leq \\text{dim}\\_{\\ell_2}(\\Psi,\\mathcal{D},1/K),$ where $\\text{dim}\\_{\\ell\\_p}$ is the $\\ell_p$ distributional eluder dimension. \nIt turns out for low-rank MDPs, we can show $\\text{dim}_{\\ell_2}(\\Psi,\\mathcal{D},\\epsilon)=\\mathcal{O}(d\\log(1/\\epsilon))$. The intuition is to leverage the linear transitions by using the one-step back trick (Lemma 12 of Uehara et al., 2021 \"Representation learning for…\") and then apply the elliptical potential lemma (Lemmas 19 and 20 of Uehara et al., 2021). Therefore, we've shown that the V-type LSEC (for any function class) is bounded by $\\mathcal{O}(d\\log(d/\\epsilon))$ in low-rank MDPs. Therefore, our framework can prove the first small-loss PAC bound for low-rank MDPs (more generally, any model with low Bellman eluder dimension!), which is significantly more general than latent variable models. We'd like to thank you again for your insightful suggestions. We hope this, along with the global response on BC, clears up any concerns about generality. \n4. **a)** As you pointed out in [6], model-free DEC can be sub-optimal in RL. However, in CBs, the DEC-like approach we took is actually rate-optimal, as shown by Section 3 of Foster and Rakhlin, 2020 “Beyond UCB: …”.  \\\n**b)** As you mentioned, another route is to use optimism with Eluder techniques, which allows for offline oracles. However, this typically has two drawbacks. First, it is more restrictive since the examples covered by it are typically linear. Second, for general function classes, the optimism step is typically computationally hard, while our CB alg is easily implementable (hence, we could provide experimental results on complex, real-world CBs).  \\\n**c)** Finally, we want to highlight that this Eluder route is captured as a special case of our online RL algorithm when $H=1$. In this one-step case, the BC assumption is not needed since there is no TD. Using our new result in (3), we can obtain a regret bound with the Eluder dimension. In sum, our algorithms in this paper can capture both the DEC and the Eluder route for CBs.\n\nThank you again for all your constructive feedback. We will incorporate all these comments in the camera-ready, and please let us know if you have any additional questions. "}]

Author rebuttal:
We are grateful for all the encouraging and constructive reviews, which have been helpful in improving and polishing our work this week. Amongst all four reviewers, three of them (9rQf, nAEh, 8JhK) inquired about the necessity and generality of distributional Bellman completeness (DistBC), so we'd like to address this in the global response.

**Necessity of BC for TD:**
As remarked in Lines 224-230, BC is necessary for TD-style algorithms to succeed. Without it, TD can diverge or converge to bad fixed points, e.g. Tsitsiklis and Van Roy, 1996 showed such a counterexample. Since our algorithms are distributional versions of GOLF and Bellman-consistent pessimism, which are TD-style algorithms that already require BC, it is quite natural for our results to rely on analogous assumptions to these prior works. One reviewer (8JhK) brought up the issue of non-monotonicity of BC (adding a new function can violate BC): we want to point out our theorems also hold under "generalized BC," a weaker and monotone assumption that there exists function classes $\mathcal{G}\_h$ such that $\mathcal{T}\_h\mathcal{F}\_{h+1}\subseteq \mathcal{G}\_h$ for all $h$ (cf. Assumption 14 of Jin et al., 2021a). If $\mathcal{G}=\mathcal{F}$, this recovers the typical BC assumption. We'll add this as a remark.

In fact, Foster et al., 2022 "Offline RL: Fundamental Barriers..." showed a lower bound that $Q^\star$-realizability and all-policy concentrability (stronger coverage condition than the single-policy one we use!) are not sufficient conditions for sample efficient offline RL. This suggests that removing BC is challenging and would require some other assumptions in its place. Seeking alternative conditions to BC is not our goal here. Instead, our contribution is that distributional versions of prior algorithms can yield small-loss bounds, which provides the first theoretical answer for the long-standing conundrum in the RL community: when optimizing expected returns, why can learning the return distribution and only then computing its mean perform better than learning the mean directly, which is all we need for Bellman's equation?



**(New result) DistBC is satisfied by linear and low-rank MDPs:**
Urged by the excellent feedback, we've proven a new result: in linear MDPs, the following linear function class automatically satisfies DistBC, $\mathcal{F}=\{f(z\mid x,a) = \phi^\star(x,a)^\top w(z): w: [0,1]\to B^d(r)\}$, where $\phi^\star(x,a)$ are the linear MDP's features and $B^d(r)$ is the radius-$r$ $\ell_2$-ball in $\mathbb{R}^d$ (with $r$ chosen appropriately for normalization purposes). This result is analogous to the well-known fact that linear MDPs automatically satisfy (vanilla) BC with a similar linear function class, and the proofs are similar. Moreover, this result easily extends to low-rank MDPs (where $\phi^\star$ is unknown) if we let $f(z\mid x,a)=\phi(x,a)^\top w(z)$ with $\phi$ varying in $\Phi$, assuming it is realizable, $\phi^\star\in\Phi$. We want to point out that DistBC also captures the Linear-Quadratic Regulator (LQR), as shown in Section B.2 of Wu et al., 2023 "Distributional Offline Policy Evaluation...". In sum, DistBC captures linear MDPs, low-rank MDPs, and LQRs, so DistBC essentially captures the same interesting models captured by vanilla BC.


**Conclusion about BC:**
To summarize the above, (1) BC is necessary for TD and assumed by GOLF and Bellman-consistent pessimism (the non-distributional analogs of our algs), and (2) DistBC captures all the interesting models captured by (vanilla) BC. We hope these two points clarify our rationale for assuming DistBC and that it is not strong at all (relative to prior works).

Additionally, we have attached a 1 page PDF containing a proof sketch of Theorem 5.2, the definition of low-rank MDP and a discussion on computational complexity, which we hope clears up any missing notations. 


Meta-review: All the reviewers are in agreement that this paper provides a strong conceptual and technical contribution to the burgeoning area of distributional RL. Based on the author-reviewer discussion, the authors are encouraged to include a detailed discussion of the distributional Bellman completeness assumption, including why it is necessary and under what conditions it is satisfied. I also suggest that the authors include the proof that it is satisfied for linear and low-rank MDPs if they are able to fit it in the camera-ready version.
Decision: Accept (poster)