Title: The Benefits of Being Distributional: Small-Loss Bounds for Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning Theory

Distributional Reinforcement Learning

Small-Loss Bounds

First-order regret
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

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 ℓ 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.

1 Introduction

The goal of reinforcement learning (RL) is to learn a policy that minimizes/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 to these estimates. By Bellman's equation, Q -functions suffice for optimal decision-making and indeed these approaches have vanishing regret bounds, suggesting we 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 can actually yield state-of-the-art performance in stratospheric balloon navigation [Bellemare et al., 2020], robotic grasping [Bodnar et al., 2020], algorithm discovery [Fawzi et al., 2022] and game playing benchmarks [Hessel et al., 2018, Dabney et al., 2018a, Barth-Maron et al., 2018]. In both online [Yang et al., 2019] and offline RL [Ma et 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.

Despite learning the whole loss distribution, DistRL algorithms use 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 a 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 et 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 an auxiliary task that leads to better representation learning. However, the theoretical

question remains: does DistRL, i.e. , learning the distribution and then computing the mean, yield provably stronger finite-sample guarantees and if so stronger how and when?

In this paper, we provide the first mathematical basis for the benefits of DistRL via the lens of small-loss bounds, which are instance-dependent bounds that depend on the minimum achievable cost in the problem [Agarwal et al., 2017]. 1 For example in linear MDPs, typical worst-case regret bounds scale on the order of poly( d, H ) √ 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 ) √ K · 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 a faster convergence rate of O (poly( d, H ) log( K )) , compared to the O (poly( d, H ) √ K ) rate in standard uniform bounds. Since we always have V ⋆ ≤ 1 , small-loss bounds simply match the standard uniform bounds in the worst case.

As warm-up, we show that maximum likelihood estimation (MLE), i.e. , maximizing log-likelihood, can be used to 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 to 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 a novel regret decomposition with triangular discrimination and also introduces the ℓ 1 distributional eluder dimension, which generalizes the ℓ 2 distributional eluder dimension of Jin et 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-loss bounds in offline RL. 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 small-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:

As warm-up, we propose a distributional CB algorithm and prove that it obtains a smallloss regret bound (Section 4). We empirically demonstrate it outperforms state-of-the-art CB algorithms in three challenging benchmark tasks (Section 7).

We propose a distributional online RL algorithm that enjoys small-loss bounds in settings with low ℓ 1 distributional eluder dimension, which we show can always capture low-rank MDPs. The ℓ 1 distributional eluder dimension may be of independent interest (Section 5).

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 a 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 guarantees 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 non-distributional algorithms. In contrast, our work shows that distributional RL yields adaptive finite-sample bounds that converge faster when the optimal cost of the problem is small. Wu et al. [2023] recently derived finite-sample bounds for distributional off-policy evaluation with MLE, while our offline RL section focuses on off-policy optimization.

First-order bounds in bandits When maximizing rewards, first-order 'small-return' bounds can be easily derived from EXP4 [Auer et al., 2002], since receiving the worst reward 0 with probability (w.p.) δ contributes at most R ⋆ δ to the regret 2 . When minimizing costs, receiving the worst loss 1 w.p. δ may induce large regret relative to L ⋆ if L ⋆ is small. To illustrate, if R ⋆ = 0 then all policies are optimal, so no learning is needed and the small-return bound is vacuous. Yet if L ⋆ = 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, 2021].

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( ∑ h r h | π k ) ≤ c · V π k ≤ c · V ⋆ . However, the 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 Lee 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.

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-Risk (CVaR), of the loss [Dabney et al., 2018b]. 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 × A → ∆([0 , 1]) . Throughout, we fix some dominating measure λ on [0 , 1] ( e.g. , Lebesgue for continuous or counting for discrete) and let ∆([0 , 1]) be all distributions on [0 , 1] that are absolutely continuous with respect to λ . We identify such a distribution with its density with respect to λ , and we also write C ( y | x, a ) for ( C ( x, a ))( y ) . Let K denote the number of episodes. At each episode k ∈ [ K ] , the learner observes a context x k ∈ X , samples an action a k ∈ A , and then receives a cost c t ∼ C ( x t , a t ) , which we assume to be normalized, i.e. , c t ∈ [0 , 1] . The goal is to design a learner that attains low regret with high probability, where regret is defined as

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 P h : X ×A → ∆( X ) and cost distributions C h : X ×A → ∆([0 , 1]) at each step h ∈ [ H ] . We start with the Online RL setting, which proceeds over K episodes as follows: at each episode k ∈ [ K ] , the learner plays a policy π k ∈ [ X → ∆( A )] H ; we start from a fixed initial state x 1 ; then for each h = 1 , 2 , . . . , H , the policy samples an action a h ∼ π k h ( x h ) , receives a cost c h ∼ C h ( x h , a h ) , and transitions to the next state x h +1 ∼ P h ( x h , a h ) . Our goal is to compete with the optimal policy that minimizes expected the loss, i.e. , π ⋆ ∈ arg min π ∈ Π V π where V π = E π [ ∑ H h =1 c h ] . Regret bounds aim to control the learner's regret with high probability, where regret is defined as,

If the algorithm returns a single policy ̂ π , it is desirable to obtain a Probably Approximately Correct (PAC) bound on the sub-optimality of ̂ π , i.e. , V ̂ π -V ⋆ .

The third setting we study is Offline RL , where instead of needing to actively explore and collect data ourselves, we are given H datasets D 1 , D 2 , . . . , D H to learn a good policy ̂ π . Each D h contains N i.i.d. samples ( x h,i , a h,i , c h,i , x ′ h,i ) from the process ( x h,i , a h,i ) ∼ ν h , c h,i ∼ C h ( x h,i , a h,i ) , x ′ h,i ∼ P h ( x h,i , a h,i ) , where ν h ∈ ∆( 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 a PAC guarantee for ̂ π , which should improve over the data generating process.

Distributional RL For a policy π and h ∈ [ H ] , let Z π h ( x h , a h ) ∈ ∆([0 , 1]) denote the distribution of the loss-to-go ∑ H t = h c t conditioned on rolling in π from x h , a h . The expectation of the above is Q π h ( x h , a h ) = ¯ Z π h ( x h , a h ) and V π h ( x h ) = E a h ∼ π h ( x h ) [ Q π h ( x h , a h )] . We use Z ⋆ h , Q ⋆ h , V ⋆ h to denote these quantities with π ⋆ . Recall the regular Bellman operator acts on a function f : X × A → [0 , 1] as follows: T π h f ( x, a ) = ¯ C h ( x, a ) + E x ′ ∼ P h ( x,a ) ,a ′ ∼ π ( x ′ ) [ f ( x ′ , a ′ )] . Analogously, the distributional Bellman operator [Morimura et al., 2012, Bellemare et al., 2017] acts on a conditional distribution d : X ×A → ∆([0 , 1]) as follows: T π,D h d ( x, a ) D = C h ( x, a ) + d ( x ′ , a ′ ) , where x ′ ∼ P h ( x, a ) , a ′ ∼ π ( x ′ ) and D = denotes equality of distributions. Another way to think about the distributional Bellman operator is that a sample z ∼ T π,D h d ( x, a ) is generated as follow: z := c + y, where c ∼ C h ( x, a ) , x ′ ∼ P h ( x, a ) , a ′ ∼ π ( x ′ ) , y ∼ d ( x ′ , a ′ ) . We will also use the Bellman optimality operator T ⋆ h and its distributional variant T ⋆,D h , defined as follows: T ⋆ h f ( x, a ) = ¯ C h ( x, a ) + E x ′ ∼ P h ( x,a ) [min a ∈A f ( x ′ , a ′ )] and T ⋆,D h d ( x, a ) D = C h ( x, a ) + d ( x ′ , a ′ ) where x ′ ∼ P h ( x, a ) , a ′ = arg min a ¯ d ( x ′ , a ) . Please see Table 2 for an index of notations.

4 Warm up: Small-Loss Regret for Distributional Contextual Bandits

̸

In this section, we propose an efficient reduction from CB to 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 a distribution over actions given predictions ̂ f ∈ R A and a parameter γ ∈ R ++ : setting b = arg min a ∈A ̂ f ( a ) as the best action with respect to the predictions, the weight for any other action a = b is,

̸

and the rest of the weight is allocated to b : ReIGW γ ( ̂ f, γ )[ b ] = 1 -∑ a = b ReIGW γ ( ̂ f, γ )[ a ] .

Algorithm 1 Distributional CB (DISTCB)

1: Input: number of episodes K , failure probability δ , ReIGW learning rate γ .

2: Initialize any cost distribution f (1) .

3: for episode k = 1 , 2 , . . . , K do

4: Observe context x k .

5: Sample action a k ∼ p k = ReIGW ( ¯ f ( k ) ( x k , · ) , γ ) from Eq. (1).

6: Observe cost c k ∼ C ( x k , a k ) and update online MLE oracle with (( x k , a k ) , c k ) .

7: end for

Wepropose Dist ributional C ontextual B andit (DISTCB) in Algorithm 1, a two-step procedure for each episode k ∈ [ K ] . Upon seeing context x k , DISTCB first samples an action a k from ReIGW generated by means of our estimated cost distributions for each action, i.e. , ̂ f ( a ) = ¯ f ( k ) ( x k , a ) , ∀ a ∈ A (Line 5). Then, DISTCB updates f ( k ) ( · | x k , a k ) by maximizing the log-likelihood to estimate the conditional cost distribution C ( · | x k , a k ) (Line 6). Formally, this second step is achieved via an online MLE oracle with a realizable distribution class F CB ⊂ X × A → ∆([0 , 1]) ; let Regret log ( K ) be some upper bound on the log-likelihood regret for all possibly adaptive sequences { x k , a k , c k } k ∈ [ K ] ,

Under realizability , C ∈ F CB , we expect Regret log ( K ) ∈ O (log( K )) . For instance, if F CB is finite, exponentially weighted average forecaster guarantees Regret log ( K ) ≤ log |F CB | [Cesa-Bianchi and Lugosi, 2006, Chapter 9]. We now state our main result for DISTCB.

The dominant term scales with the optimal sum of costs √ 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 smallloss regret [Foster and Krishnamurthy, 2021, Theorem 1]. Our bound matches that of FastCB in terms of dependence on A,C ⋆ and log(1 /δ ) . Our key difference with FastCB is the online supervised learning oracle: in DISTCB, we aim to learn the conditional cost distribution by maximizing loglikelihood, while FastCB aims to perform regression with the binary cross-entropy loss. In Section 7, we find that DISTCB empirically outperforms SquareCB and FastCB in 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,

For any distributions f, g ∈ ∆([0 , 1]) , their triangular discrimination 3 is defined as D △ ( f ∥ g ) := ∫ ( f ( y ) -g ( y )) 2 f ( y )+ g ( y ) d λ ( y ) . The key insight is that ⋆ can be bounded by the triangular discrimination of f ( k ) ( s k , a k ) and C ( s k , a k ) : by Cauchy-Schwartz and y 2 ≤ y for y ∈ [0 , 1 ], we have ¯ f -¯ g = ∫ y ( f ( y ) -g ( y ))d λ ( y ) ≤ √ ∫ y ( f ( y ) + g ( y ))d λ ( y ) √ ∫ ( f ( y ) -g ( y )) 2 f ( y )+ g ( y ) d λ ( y ) , and hence,

So, Eq. ( △ 1 ) implies that ⋆ is bounded by D △ ( f ( k ) ( s k , a k ) ∥ C ( s k , a k )) . Since D △ is equivalent (up to universal constants) to the squared Hellinger distance, Foster et al. [2021, Lemma A.14] implies the above can be bounded by the online MLE regret, so w.p. at least 1 -δ , we have

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 to approximate MLE. To leverage function approximation for learning conditional distributions, we use a generic function class F ⊆ ( X × A → ∆([0 , 1])) H where each element f ∈ F is a tuple f = ( f 1 , . . . , f H ) such that each f h is a candidate estimator for Z ⋆ h , the distribution of loss-to-go ∑ H t = h c t under π ⋆ . For notation, f H +1 ( x, a ) = δ 0 denotes the dirac at zero for all x, a .

We now present our O ptimistic Dis tributional C onfidence set O ptimization (O-DISCO) algorithm in Algorithm 2, consisting of three key steps per episode. At episode k ∈ [ K ] , O-DISCO first identifies the f ( k ) with the minimal expected value at h = 1 over the previous confidence set F k -1 (Line 4). This step induces global optimism . Then, O-DISCO collects data for this episode by rolling in with the greedy policy π k with respect to the mean of f ( k ) (Line 6). Finally, O-DISCO constructs a

Algorithm 2 O ptimistic Dis tributional C onfidence set O ptimization (O-DISCO)

1: Input: number of episodes K , distribution class F , threshold β .

2: Initialize D h, 0 ←∅ for all h ∈ [ H ] , and set F 0 = F .

3: for episode k = 1 , 2 , . . . , K do

4: Set optimistic estimate f ( k ) = arg min f ∈F k -1 min a ¯ f 1 ( x 1 , a ) .

5: Set π k h ( x ) = arg min a ¯ f ( k ) h ( x, a ) .

6: Roll out π k and obtain a trajectory x 1 ,k , a 1 ,k , c 1 ,k , . . . , x H,k , a H,k , c H,k . For each h ∈ [ H ] , augment the dataset D h,k = D h,k -1 ∪ { ( x h,k , a h,k , c h,k , x h +1 ,k ) } .

7: For all ( h, f ) ∈ [ H ] ×F , sample y f h,i ∼ f h +1 ( x ′ h,i , a ′ ) and a ′ = arg min a ¯ f h +1 ( x ′ h,i , a ) , where ( x h,i , a h,i , c h,i , x ′ h,i ) is the i -th datapoint of D h,k . Then, set z f h,i = c h,i + y f h,i and define the confidence set

8: end for

9: Output: ¯ π = unif( π 1: K ) .

confidence set F k by including a function f if it exceeds a threshold on the log-likelihood objective using data z f h,i ∼ T ⋆,D h f h +1 ( x h,i , a h,i ) for all steps h simultaneously (Line 7). This step is called local fitting , as each f ∈ F k has the property that f h is close-in-distribution to T ⋆,D h f h +1 for all h . We highlight that O-DISCO only learns the distribution for estimating the mean, i.e. , Lines 4 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.

To ensure that MLE succeeds for the Temporal-Difference (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 π, h ∈ [ H ] , f h +1 ∈ F h +1 = ⇒ T π,D h f h +1 ∈ F h .

5.1 The ℓ 1 Distributional Eluder Dimension

We now introduce the ℓ 1 distributional eluder dimension. Let S be an abstract input space, let Ψ be a set of functions mapping S → R and let D be a set of distributions on S .

Definition 5.2 ( ℓ p -distributional eluder dimension) . For any function class Ψ ⊆ S → R , distribution class D ⊆ ∆( S ) and ε > 0 , the ℓ p -distributional eluder dimension (denoted by DE p (Ψ , D , ε ) ) is the length L of the longest sequence d (1) , d (2) , . . . , d ( L ) ⊆ D such that there exists ε ′ ≥ ε , such that for all t ∈ [ L ] , we have that there exists f ∈ Ψ such that | E d ( t ) f | > ε and also ∑ t -1 i =1 | E d ( i ) f | p ≤ ε p .

When p = 2 , this is exactly the ℓ 2 distributional eluder of Jin et al. [2021a, Definition 7]. We're particularly interested in the p = 1 case, which can be used with MLE's generalization bounds. The following is a key pigeonhole principle for the ℓ 1 distributional eluder dimension.

Theorem 5.3. Let C := sup d ∈D ,f ∈ Ψ | E d f | be the envelope. Fix any K ∈ N and sequences f (1) , . . . , f ( K ) ⊆ Ψ , d (1) , . . . , d ( K ) ⊆ D . Let β be a constant such that for all k ∈ [ K ] , we have, ∑ k -1 i =1 ∣ ∣ E d ( i ) f ( k ) ∣ ∣ ≤ β. Then, for all k ∈ [ K ] , we have

As we'll see later, Theorem 5.3 is the key tool that transfers triangular discrimination guarantees on the training distribution to any new test distribution. Another key property is that the ℓ 1 dimension generalizes the original ℓ 2 dimension of Jin et al. [2021a].

Finally, we note that our distributional eluder dimension generalize the regular ℓ 1 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 × 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 π k for a single trajectory, partially roll-out π k for H times where for each h ∈ [ H ] , we collect x h,k ∼ d π k h , take a random action a h,k ∼ unif( A ) , observe c h,k ∼ C h ( x h,k , a h,k ) , x ′ h,k ∼ P h ( x h,k , a h,k ) and augment the dataset D h,k = D h,k -1 ∪ { ( x h,k , a h,k , c h,k , x ′ h,k ) } . 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 D h = { ( x, a ) ↦→ d π h ( x, a ) : π ∈ Π } and Ψ h = { ( x, a ) ↦→ D △ ( f ( x, a ) ∥ T ⋆,D f ( x, a )) : f ∈ F } . Similarly, the V-type classes are D h,v = { x ↦→ d π h ( x ) : π ∈ Π } and Φ h,v = { x ↦→ E a ∼ Unif( A ) [ D △ ( f ( x, a ) ∥ T ⋆,D f ( x, a ))] : f ∈ F } . Finally, define DE 1 ( ε ) = max h DE 1 (Ψ h , D h , ε ) and DE 1 ,v ( ε ) = max h DE 1 (Ψ h,v , D h,v , ε ) .

Theorem 5.5. Suppose DistBC holds (Assumption 5.1). For any δ ∈ (0 , 1) , w.p. at least 1 -δ , running O-DISCO with β = log( HK |F| /δ ) guarantees the following regret bound,

Regret O-DISCO ( K ) ≤ 160 H √ KV ⋆ DE 1 (1 /K ) log( K ) β +18000 H 2 DE 1 (1 /K ) log( K ) β. If UAE = TRUE (Algorithm 4), then the learned mixture policy ¯ π is guaranteed to satisfy,

Compared to prior bounds for GOLF [Jin et al., 2021a], the leading √ K 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, if V ⋆ ≤ O (1 / √ K ) , then our regret bound converges at a fast O ( H 2 DE 1 (1 /K ) log( K ) β ) rate. While there are existing first-order bounds in online RL, 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 4 . 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 FittedQ -Evaluation (FQE) can diverge or converge to bad fixed point solutions [Tsitsiklis and Van Roy, 1996, Munos and Szepesvári, 2008, Kolter, 2011]. As a 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 G h such that f h +1 ∈ F h +1 = ⇒ T π,D h f h +1 ∈ G h [as in Jin et al., 2021a, Assumption 14]. Simply replace max g ∈F h in the confidence set construction with max g ∈G h . While adding functions to F may break BC (as BC is not monotonic), we can always augment G to satisfy generalized completeness.

Computational complexity When taken as is, OLIVE [Jiang et al., 2017], GOLF, and our algorithms are version space methods that suffer from a 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, ε -greedy suffices in problems that don't require deep and strategic exploration, i.e. , a large myopic exploration gap [Dann et al., 2022]. With ε -greedy, a replay buffer, and discretization, our algorithm essentially recovers C51 [Bellemare et al., 2017]. We leave developing and analyzing computationally efficient algorithms based on our insights as promising future work.

5.3 Instantiation with Low-Rank MDPs

The low-rank MDP [Agarwal et al., 2020] is a standard abstraction for non-linear function approximation used in theory [Uehara et al., 2021] and practice [Zhang et al., 2022, Chang et al., 2022].

Definition 5.6 (Low-rank MDP) . A transition model P h : X × A → ∆( X ) has rank d if there exist unknown features ϕ ⋆ h : X × A → R d , µ ⋆ h : X → R d such that P h ( x ′ | x, a ) = ϕ ⋆ h ( x, a ) ⊤ µ ⋆ h ( x ′ ) for all x, a, x ′ . Also, assume max x,a ∥ ϕ ⋆ h ( x, a ) ∥ 2 ≤ 1 and ∥ ∫ g d µ ⋆ h ∥ 2 ≤ ∥ g ∥ ∞ √ d for all functions g : X → R . The MDP is called low-rank if P h is low-rank for all h ∈ [ H ] .

We now specialize Theorem 5.5 to low-rank MDPs with three key steps. First, we bound the V-type eluder dimension by DE 1 ,v ( ε ) ≤ O ( d log( d/ε )) , which is a known result that we reproduce in Theorem G.4. The next step requires access to a realizable Φ class, i.e. , for all h ∈ [ H ] , ϕ ⋆ h ∈ Φ , which is a standard assumption for low-rank MDPs [Agarwal et al., 2020, Uehara et al., 2021, Mhammedi et al., 2023]. Given the realizable Φ , we can construct a specialized F for the low-rank MDP: F lin = F lin 1 ×··· × F lin H ×F lin H +1 where F lin H +1 = { δ 0 } and for all h ∈ [ H ] ,

where α := max h,π,z,x,a Z π h ( z | x, a ) is the largest mass for the cost-to-go distributions. In Appendix D, we show that F lin satisfies DistBC. Further, if costs are discretized into a uniform grid of M points, its bracketing entropy is bounded by ˜ O ( dM +log | Φ | ) . Discretization is necessary to bound the statistical complexity of F lin 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 δ ∈ (0 , 1) , w.p. at least 1 -δ , running O-DISCO with UAE = TRUE and with F lin as described above learns a policy ¯ π such that,

Proof. As described above, we have DE 1 (1 /K ) ≤ O ( d log( dK )) and β = log( HK/δ ) + dM + log | Φ | . Since DistBC is satisfied by F lin , 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 ⋆ ≤ ˜ O (1 /K ) , O-DISCO has a fast ˜ O (1 /K ) convergence rate which improves over all prior results that converge at a slow ˜ Ω(1 / √ K ) rate [Uehara et 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 k : (i) Z ⋆ ∈ F k , and (ii) elements of F k almost satisfy the distributional Bellman equation, i.e. , for all h ∈ [ H ] , we have ∑ k i =1 E π i [ δ h,k ( x h , a h )] ≤ O ( β ) where δ h,k ( x h , a h ) = D △ ( f ( k ) h ( x h , a h ) ∥ T ⋆,D h f ( k ) h +1 ( x h , a h )) . Next, we derive a corollary of Eq. ( △ 1 ):

To see why this is true, apply AM-GM to Eq. ( △ 1 ) to get 2( ¯ f -¯ g ) ≤ ¯ f + ¯ g + D △ ( f ∥ g ) , which simplifies to ¯ f ≤ 3¯ g + D △ ( f ∥ g ) . Plugging this back into Eq. ( △ 1 ) yields Eq. ( △ 2 ). Then, by iterating Eq. ( △ 2 ) and AM-GM, we derive a self-bounding lemma: for any f, π, h , we have ¯ f h ( x h , a h ) ≲ Q π h ( x h , a h ) + H ∑ H t = h E π,x h ,a h [ D △ ( f t ( x t , a t ) ∥ T π,D h f t +1 ( x t , a t ))] (Lemma H.3).

The implicit inequality V π k -V ⋆ ≲ √ V ⋆ + H ∑ H h =1 E π k [ δ h,k ( x h , a h )] √ H E π k [ δ h,k ( x h , a h )] can then be obtained by AM-GM and rearranging. The final step is to sum over k and bound ∑ K k =1 E π k [ δ h,k ( x h , a h )] via the eluder dimension's pigeonhole principle (Theorem 5.3 applied with fact (ii)). Please see Appendix H for the full proof.

6 Small-Loss Bounds for Offline Distributional RL

We now propose P essimistic Dis tributional C onfidence set O ptimization (P-DISCO; Algorithm 3), which adapts the distributional confidence set technique from the previous section to the offline setting by leveraging pessimism instead of optimism. Notably, P-DISCO is a 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 a similar log-likelihood thresholding procedure as in O-DISCO, where the difference is we now use data sampled from T π,D h f h +1 instead of T ⋆,D h f h +1 . Next, output the policy with the most pessimistic mean amongst all the confidence sets.

Algorithm 3 P essimistic Dis tributional C onfidence set O ptimization (P-DISCO)

1: Input: datasets D 1 , . . . , D H , distribution function class F , threshold β , policy class Π .

2: For all ( h, f, π ) ∈ [ H ] × F × Π , sample y f,π h,i ∼ f h +1 ( x ′ h,i , π h +1 ( x ′ h,i )) , where ( x h,i , a h,i , c h,i , x ′ h,i ) is the i -th datapoint of D h . Then, set z f,π h,i = c h,i + y f,π h,i and define the confidence set,

3: For each π ∈ Π , define the pessimistic estimate f π = arg max f ∈F π E a ∼ π ( x 1 ) [ ¯ f 1 ( x 1 , a ) ] .

4: Output: ̂ π = arg max π ∈ Π E a ∼ π ( x 1 ) [ ¯ f π 1 ( x 1 , π ) ] .

In offline RL, many works made strong all-policy coverage assumptions [Antos et al., 2008, Chen and Jiang, 2019]. Recent advancements [Kidambi et al., 2020, Xie et al., 2021, Uehara and Sun, 2022, Rashidinejad et al., 2021, Jin et al., 2021b] have pursued best effort guarantees that aim to compete with any covered policy ˜ π , with sub-optimality of the learned ̂ π degrading gracefully as coverage worsens. The coverage is measured by the single-policy concentrability C ˜ π = max h ∥ ∥ d d ˜ π h / d ν h ∥ ∥ ∞ . 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 δ ∈ (0 , 1) , w.p. at least 1 -δ , running P-DISCO with β = log( H | Π ||F| /δ ) learns a policy ̂ π that enjoys the following PAC bound with respect to any comparator policy ˜ π ∈ Π :

To the best of our knowledge, this is the first small-loss bound for offline RL, 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 a comparator policy ˜ π if one of the following is true: (i) ν has good coverage over ˜ π , so the O (1 / √ N ) term is manageable; or (ii) ˜ π has small-loss, in which case we may even obtain a fast O (1 /N ) rate. Thus, P-DISCO has two chances at strongly competing with ˜ π , while conventional offline RL methods solely rely on (i) to be true.

7 Distributional CB Experiments

We now compare DISTCB with SquareCB [Foster and Rakhlin, 2020] and the state-ofthe-art CB method FastCB [Foster and Krishnamurthy, 2021], which respectively minimize the squared loss and log loss for estimating the conditional mean. The key question we investigate here is whether learning the conditional mean 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 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 is 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 is . 1 × ( y -ˆ y ) where y is the actual risk level, and ˆ y is the predicted risk level.

CIFAR-100 This popular image dataset contains 100 classes, which correspond to our actions, and each class is in one of 20 superclasses. We assign cost as 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.

Results Across tasks, DISTCB achieves lower average cost over 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-entropy). An even more fine-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://github.com/kevinzhou497/distcb .

8 Conclusion

We showed that distributional RL leads to small-loss bounds in both online and offline RL, and we also proposed a distributional CB algorithm that outperforms the state-of-the-art FastCB. A fruitful direction would be to investigate connections of natural policy gradient with our MLE distributionalfitting 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 or even faster bounds.

Acknowledgements This material is based upon work supported by the National Science Foundation under Grant Nos. IIS-1846210 and IIS-2154711.

References

Alekh Agarwal, Akshay Krishnamurthy, John Langford, Haipeng Luo, et al. Open problem: Firstorder regret bounds for contextual bandits. In Conference on 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-20107, 2020.

Alekh Agarwal, Yuda Song, Wen Sun, Kaiwen Wang, Mengdi Wang, and Xuezhou Zhang. Provable 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 bellmanresidual 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-77, 2002.

Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement 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://openreview.net/ forum?id=SyZipzbCb .

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, Subhodeep Moitra, Sameera S Ponda, and Ziyu Wang. Autonomous navigation of stratospheric balloons using reinforcement learning. Nature , 588(7836):77-82, 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 Cesa-Bianchi and Gábor Lugosi. Prediction, learning, and games . Cambridge university press, 2006.

Jonathan Chang, Kaiwen Wang, Nathan Kallus, and Wen Sun. Learning bellman complete representations for offline policy evaluation. In International Conference on Machine Learning , pages 2938-2971. PMLR, 2022.

Jinglin Chen and Nan Jiang. Information-theoretic considerations in batch reinforcement learning. In International Conference on 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, Marc Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 32, 2018b.

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 preprint arXiv:2210.13573 , 2022.

Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, Mohammadamin Barekatain, Alexander Novikov, Francisco J R Ruiz, Julian Schrittwieser, Grzegorz Swirszcz, 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 J 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-3489. PMLR, 2022.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , pages 770-778, 2016. doi: 10.1109/CVPR.2016.90.

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. Contextual 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-4879. PMLR, 2020a.

Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I 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 in neural information processing systems , 34:13406-13418, 2021a.

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? 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 , 33:15312-15325, 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: Modelbased 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-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 in neural information processing systems , 33:15522-15533, 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-2218, 2022.

Clare Lyle, Marc 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 in Neural Information Processing Systems , 34:19235-19247, 2021.

Zakaria Mhammedi, Adam Block, Dylan J 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 A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 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://kaggle.com/competitions/prudential-life-insurance-assessment .

Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka. Parametric return density estimation for reinforcement learning. arXiv preprint arXiv: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.

Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline reinforcement learning and imitation learning: A tale of pessimism. Advances in Neural Information Processing Systems , 34:11702-11716, 2021.

Mark Rowland, Marc 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, Marc G Bellemare, and Will Dabney. An analysis of quantile temporaldifference learning. arXiv preprint arXiv:2301.04462 , 2023.

Daniel Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic exploration. Advances in Neural Information Processing Systems , 26, 2013.

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

John Tsitsiklis and Benjamin Van Roy. Analysis of temporal-diffference 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.net/forum?id=tyrJsbKAe6 .

Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. In International Conference on Learning Representations , 2021.

Sara van de Geer. Empirical Processes in M-estimation , volume 6. 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.1145/2641190.2641198. URL http://doi.acm.org/10.1145/2641190.2641198 .

István Vincze. On the concept and measure of information contained in an observation. In Contributions 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-optimal risk-sensitive reinforcement learning with cvar. International Conference on Machine Learning , 2023.

Ruosong Wang, Dean Foster, and Sham M. Kakade. What are the statistical limits of offline RL with linear function approximation? In International Conference on Learning Representations , 2021a. URL https://openreview.net/forum?id=30EvkP2aQLD .

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. PMLR, 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://openreview.net/forum?id=LQIjzPdDt3q .

Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie-Yan Liu. Fully parameterized quantile function for distributional reinforcement learning. Advances in neural information processing systems , 32, 2019.

Andrea Zanette and Emma Brunskill. Tighter problem-dependent regret bounds in reinforcement learning without domain knowledge using value function bounds. In International Conference on Machine Learning , pages 7304-7312. PMLR, 2019.

Andrea Zanette, Martin J Wainwright, and Emma Brunskill. Provable benefits of actor-critic methods for offline reinforcement learning. Advances in neural information processing systems , 34:1362613640, 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-ml.org/lt-book.html .

Zhengyuan Zhou, Susan Athey, and Stefan Wager. Offline multi-action policy learning: Generalization and optimization. Operations Research , 71(1):148-183, 2023.

A Notations

Appendices

S , A , A State and action spaces, and A = |A| . ∆( S ) The set of distributions supported by S . ¯ d The expectation of any real-valued distribution d , i.e. , ¯ d = E y ∼ d [ y ] . [ N ] { 1 , 2 , . . . , N } for any natural number N . Z π h ( x, a ) Distribution of ∑ H t = h c t given x h = x, a h = a rolling in from π . Q π h ( x, a ) , V π h ( x ) Q π h ( x, a ) = ¯ Z π h ( x, a ) and V π h = E a ∼ π ( x ) [ Q π h ( x, a )] . π ⋆ Optimal policy, i.e. , π ⋆ = arg min π V π 1 ( x 1 ) . Without loss of optimality, we take π ⋆ : X → A to be Markov & deterministic. Z ⋆ h , Q ⋆ h , V ⋆ h Z π h , Q π h , V π h with π = π ⋆ , the optimal policy. T π h , T ⋆ h The Bellman operators that act on functions. T π,D h , T ⋆,D h The distributional Bellman operators that act on conditional distributions. V π , Z π , V ⋆ , Z ⋆ V π = V π 1 ( x 1 ) , Z π = Z π 1 ( x 1 ) . V ⋆ , Z ⋆ are defined similarly with π ⋆ . d π h ( x, a ) The probability of π visiting ( x, a ) at time h . C ˜ π Coverage coefficient max h ∥ ∥ d d ˜ π h / d ν h ∥ ∥ ∞ . D △ ( f ∥ g ) Triangular discrimination between f, g . H ( f ∥ g ) Hellinger distance between f, g . D KL ( f ∥ g ) KL divergence between f, g .

A.1 Statistical Distances

Let f, g be distributions over Y . Then,

The following standard inequalities will be helpful:

Lemma A.1. For any distributions f, g , we have 2 H 2 ( f ∥ g ) ≤ D △ ( f ∥ g ) ≤ 4 H 2 ( f ∥ g ) .

Proof. Recall that

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-maximizing setting (instead of the cost-minimizing setting studied throughout the paper); if SMALLRETURN is 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 β , flag UAE, flag SMALLRETURN.

2: Initialize D h, 0 ←∅ for all h ∈ [ H ] , and set F 0 = F .

3: Set op = max if SMALLRETURN else op = min .

4: for episode k = 1 , 2 , . . . , K do

5: Set f ( k ) = arg op f ∈F k -1 op a ¯ f 1 ( x 1 , a ) .

6: Set π k h ( x ) = arg op a ¯ f ( k ) h ( x, a ) .

8: For each h ∈ [ H ] , collect x h,k ∼ d π k h , a h,k ∼ unif( A ) , c h,k ∼ C h ( x h,k , a h,k ) , x ′ h,k ∼ P h ( x h,k , a h,k ) , and augment the dataset D h,k = D h,k -1 ∪ { ( x h,k , a h,k , c h,k , x ′ h,k ) } .

7: if UAE then

9: else

10: Roll out π k and obtain a trajectory x 1 ,k , a 1 ,k , c 1 ,k , . . . , x H,k , a H,k , c H,k . For each h ∈ [ H ] , augment the dataset D h,k = D h,k -1 ∪ { ( x h,k , a h,k , c h,k , x h +1 ,k ) } .

11: end if

12: For all ( h, f ) ∈ [ H ] × F , sample y f h,i ∼ f h +1 ( x ′ h,i , a ′ ) and a ′ = arg op a ¯ f h +1 ( x ′ h,i , a ) , where ( x h,i , a h,i , c h,i , x ′ h,i ) is the i -th datapoint of D h,k . Also, set z f h,i = c h,i + y f h,i and define the confidence set,

13: end for

14: Output: ¯ π = unif( π 1: K ) .

Algorithm 5 P-DISCO (with small return)

1: Input: datasets D 1 , . . . , D H , distribution function class F , threshold β , policy class Π , flag SMALLRETURN.

2: For all ( h, f, π ) ∈ [ H ] × F × Π , sample y f,π h,i ∼ f h +1 ( x ′ h,i , π h +1 ( x ′ h,i )) , where ( x h,i , a h,i , c h,i , x ′ h,i ) is the i -th datapoint of D h . Then, set z f,π h,i = c h,i + y f,π h,i and define the confidence set,

3: Set op = max if SMALLRETURN else op = min .

5: Output: ̂ π = arg op π ∈ Π E a ∼ π ( x 1 ) [ ¯ f π 1 ( x 1 , π ) ] .

4: For each π ∈ Π , define the pessimistic estimate f π = arg op f ∈F π E a ∼ π ( x 1 ) [ ¯ f 1 ( x 1 , a ) ] .

C Proofs for DISTCB

LemmaC.1 (Azuma) . Let { X i } i ∈ [ N ] be a sequence of random variables supported on [0 , 1] , adapted to filtration {F i } i ∈ [ N ] . For any δ ∈ (0 , 1) , we have w.p. at least 1 -δ ,

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 1 1 -exp( -λ ) ≤ 2 .

Proof of Theorem 4.1. First, recall the per-step inequality of ReIGW Foster and Krishnamurthy [2021, Theorem 4], which states: for any ̂ f and γ ≥ 2 A , if we set p = ReIGW γ ( ̂ f, γ ) , then, for all f ∈ [0 , 1] A , we have

where a ⋆ = arg min a f ( a ) . For any k ∈ [ K ] , applying this to ̂ f = ¯ f ( k ) ( s k , · ) , p = p k and f = ¯ C ( s k , · ) , we have

Since D △ ≤ 4 H 2 , we have

Hence, we have

Finally, recalling that 1 / (1 -ε ) ≤ 1 + 2 ε when ε ≤ 1 2 , and the fact that 5 A γ ≤ 1 2 , we have

By Azuma's inequality, we have

(Multiplicative Azuma)

Case 2 is when the left term dominates, then the bound is,

Putting these two cases together, we have the result.

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) . A transition model P h : X × A → ∆( X ) has rank d if there exist unknown features ϕ ⋆ h : X × A → R d , µ ⋆ h : X → R d such that P h ( x ′ | x, a ) = ϕ ⋆ h ( x, a ) ⊤ µ ⋆ h ( x ′ ) for all x, a, x ′ . Also, assume max x,a ∥ ϕ ⋆ h ( x, a ) ∥ 2 ≤ 1 and ∥ ∫ g d µ ⋆ h ∥ 2 ≤ ∥ g ∥ ∞ √ d for all functions g : X → R . The MDP is called low-rank if P h is low-rank for all h ∈ [ H ] .

Suppose that we have a function class Φ such that ϕ ⋆ h ∈ Φ for all h , i.e. , Φ is a realizable function class. For example, in linear MDPs, this is automatically satisfied since we know ϕ ⋆ a priori , so Φ is the singleton with ϕ ⋆ . Having a realizable Φ class is standard for solving low-rank MDPs [Uehara et al., 2021, Agarwal et al., 2023].

In what follows, let α = max h,π,z,x,a Z π h ( z | x, a ) denote the maximum density/mass value of the loss-to-go distributions. Note that α ≥ 1 always since the mass at H +1 is deterministically placed at zero. If we further know that Z π h is discretely distributed, then α = 1 . If Z π h is continuously distributed, we assume it is bounded.

We consider the function class in Eq. (2), which we reproduce here:

The next lemma (Lemma D.1) shows that this function class satisfies distributional BC.

Lemma D.1. F lin satisfies distributional BC (Assumption 5.1).

Proof. We denote ∥ f ∥ ∞ = max z,x,a f ( z | x, a ) . For any f h +1 ∈ F lin h +1 , we have ∥ f h +1 ∥ ∞ ≤ α by the construction of F lin h +1 . Then, let T D be either the distributional Bellman operator or distributional optimality operator, the following equalities hold for the appropriate a ′ ( x ′ ) based on T D ,

Since ∫ R Pr h ( c | x, a ) f h +1 ( z -c | x ′ , a ′ ( x ′ )) d c ≤ ∥ f h +1 ∥ ∞ , we know that

We further note that

Also note that ϕ ⋆ h ∈ Φ by realizability. Therefore, T D f h +1 ∈ F lin h , which is the distributional BC condition.

D.1 Bounding the bracketing number via discretized rewards

We now bound the bracketing number of F lin h under a discretization assumption that costs and costs-to-gos can only take M many discrete values on an evenly spaced grid . This can be interpreted as discretizing the reward space, and it can be shown that this discretization error is small for regret or PAC bounds [Wang et al., 2023, Section 6]. Structural assumptions are necessary to bound the

complexity of F lin h and such discretization assumptions are common in practice, e.g. , C51 [Bellemare et al., 2017] and Rainbow [Hessel et al., 2018] both set 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 Z π h are discrete, we have α = 1 .

Now, let ε > 0 be arbitrary and fixed. Recall that the ℓ ∞ bracketing number is equivalent (up to universal constants) to the ℓ ∞ covering number, so we will work with the latter. Let B ( r ) denote the d -dimensional ball of radius r (in ℓ 2 ). Recall that the ε -covering number (in ℓ 2 ) of functions [ M ] ↦→ B ( r ) scales as O (( r/ε ) dM ) . Let W ε be such the smallest cover. We can build a ℓ ∞ cover of F lin h as follows: C ε = { ( x, a, z ) ↦→⟨ ϕ ( x, a ) , w ( z ) ⟩ , w ∈ W ε , ϕ ∈ Φ } .

To check this is a ε cover, consider any f ∈ F lin h . f corresponds to some ϕ and w . Let w ′ be the neighbor of w in W ε and let f ′ ( x, a, z ) = ⟨ ϕ ( x, a ) , w ′ ( z ) ⟩ so indeed f ′ ∈ C ε . Then, for any x, a, z , we have |⟨ ϕ ( x, a ) , w ( z ) -w ′ ( z ) ⟩| ≤ ∥ ϕ ( x, a ) ∥ 2 ∥ w ( z ) -w ′ ( z ) ∥ 2 ≤ ε . Hence, C ε is an ℓ ∞ cover of size O (( √ d/ε ) dM · | Φ | ) , and so we have shown that log N [] ( ε, F lin h , ∥ · ∥ ∞ ) ≤ O ( dM log( d/ε ) + log | Φ | ) .

Linear MDPs: Recall that in linear MDPs, we know the true ϕ ⋆ and so | Φ | = 1 . Thus, the bracketing number is simply O ( dM log( d/ε )) 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.

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 = { ( x i , y i ) } i ∈ [ n ] from a martingale process: for i = 1 , 2 , ..., n , sample x i ∼ D i ( x 1: i -1 , y 1: i -1 ) and y i ∼ p ( · | x i ) . Let f ⋆ ( x, y ) = p ( y | x ) and we are given a realizable, i.e. , f ⋆ ∈ F , function class F : X ×Y → ∆( R ) of distributions. The MLE is an estimate for f ⋆ that maximizes the log-likelihood objective over our dataset:

For our guarantees to hold for general hypotheses classes F , we use the bracketing number to quantify the statistical complexity of F [van de Geer, 2000].

Definition E.1 (Bracketing Number) . Let G be a set of functions mapping X → R . Given two functions l, u such that l ( x ) ≤ u ( x ) for all x ∈ X , the bracket [ l, u ] is the set of functions g ∈ G such that l ( x ) ≤ g ( x ) ≤ u ( x ) for all x ∈ X . We call [ l, u ] an ε -bracket if ∥ u -l ∥ ≤ ε . Then, the ε -bracketing number of G with respect to ∥·∥ , denoted by N [] ( ε, G , ∥·∥ ) is the minimum number of ε -brackets needed to cover G .

Since the triangular discrimination is equivalent to squared Hellinger up to universal constants, we now prove MLE generalization bounds in terms of squared Hellinger.

Lemma E.2. Let f 1 : X → ∆( Y ) and f 2 : X × Y → R + satisfying sup x ∈X ∫ Y f 2 ( x, y )d y ≤ s , then for any distribution D ∈ ∆( X ) , we have

Proof. This follows from the proof of Wu et al. [2023, Lemma C.1].

Lemma E.3. Fix δ ∈ (0 , 1) . Then w.p. at least 1 -δ , for any f ∈ F , we have

Rearranging, we also have

Proof. We take an ϵ -bracketing of F , { [ l i , u i ] : i = 1 , 2 , . . . } , and denote ˜ F = { u i : i = 1 , 2 , . . . } . Applying Lemma 24 of Agarwal et al. [2020] to function class ˜ F and using Chernoff method, w.p. at least 1 -δ , for all ˜ f ∈ ˜ F , we have

Now, fix any f ∈ F and pick ˜ f ∈ ˜ F as the upper bracket, i.e. , f ≤ ˜ f . Now set L ( f, D ) = ∑ n i =1 -1 / 2 log( f ⋆ ( x i , y i ) /f ( x i , y i )) . Then the right hand side of (6) is

On the other hand, since H is a metric, we have

For (iii) , by the definition, we have ˜ f ( x, y ) -f ( x, y ) ∈ [0 , ϵ ] for all x , so

For (iv) , we apply Lemma E.2 with f 1 = f ⋆ and f 2 = ˜ f (thus s = 1 + ϵ |Y| ) and get

By plugging (iii) and (iv) back we get

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 F is finite. Fix any δ ∈ (0 , 1) , set β = log( |F| /δ ) and define

Then w.p. at least 1 -δ , the following holds:

The true distribution is in the version space, i.e., f ⋆ ∈ ̂ F .

Any function in the version space is close to the ground truth data-generating distribution, i.e., for all f ∈ ̂ F

Proof. These two claims follow from Lemma E.3 with ϵ = 0 , and so N [] ( ϵ, F , ∥ · ∥ ∞ ) = |F| . For (1), apply Eq. (5) to f = ̂ f MLE to see that f ⋆ ∈ ̂ F . For (2), apply Eq. (4) and note that the sum term is at most 4 β . Thus, the right hand side of Eq. (4) is at most (6 + 8 + 8) β = 22 β .

We now state the result for infinite F using bracketing entropy.

Theorem E.5. Fix any δ ∈ (0 , 1) , set β = log( N [] (( n |Y| ) -1 , F , ∥ · ∥ ∞ ) /δ ) and define

Then w.p. at least 1 -δ , the following holds:

The true distribution is in the version space, i.e., f ⋆ ∈ ̂ F .

Any function in the version space is close to the ground truth data-generating distribution, i.e., for all f ∈ ̂ F

Proof. These two claims follow from Lemma E.3 with ϵ = 1 / n |Y| . For (1), apply Eq. (5) to f = ̂ f MLE to see that f ⋆ ∈ ̂ F . For (2), apply Eq. (4) and note that the sum term is at most 7 β . Thus, the right hand side of Eq. (5) is at most (6 + 14 + 8) β = 28 β .

F Confidence set construction with general function class

In this section, we 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 on an ε -net of F . While constructing an ε -net for F is admittedly a computationally hard procedure, this is still information theoretically possible and our focus in O-DISCO and P-DISCO is to show that distributional RL information-theoretically leads to small-loss bounds.

We first define some notations. Let F ↓ and F ↑ denote a lower and upper ε -bracketing of F , i.e. , for any f ∈ F , there exists an ε -bracket [ f ↓ , f ↑ ] such that for all h , f ↓ h ≤ f h ≤ f ↑ h with f ↓ ∈ F ↓ , f ↑ ∈ F ↑ . Recall that a lower bracket g ∈ F ↓ may not be a valid distribution, but since elements of F map to non-negative values, we can assume g has non-negative entires as well. Also, we have α g h ( x, a ) := ∫ g h ( z | x, a ) ≥ 1 -ε , so for ε small enough, g is normalizable. Hence, define ˜ g ( z | x, a ) = α g h ( x, a ) -1 g ( z | x, a ) as the normalized version, which is a valid distribution that we can sample from.

Now, consider any martingale { x h,i , a h,i , c h,i } i ∈ [ n ] ,h ∈ [ H ] , which could be the online data up to episode k or the offline data (consisting of N i.i.d. samples). We define the MLE with respect to a lower bracket element as follows. For any h ∈ [ H ] , g ∈ F ↓ , π ∈ Π , sample y g,π h,i ∼ ˜ g h +1 ( x ′ h,i , π ( x ′ h,i )) , and z g,π h,i = c h,i + y g,π h,i , define the MLE solution for ( g, π ) at time h as,

Also, define the version space with respect to the above MLE as,

We now prove a key result that implies that T π h f ↓ h +1 falls into the confidence set F f ↓ ,π,h .

Theorem F.1. For any δ ∈ (0 , 1) and suppose n ≥ 2 . Then, w.p. at least 1 -δ , for any h ∈ [ H ] , g ∈ F , f ↓ ∈ F ↓ , π ∈ Π , we have

Proof of Theorem F.1. Consider a ε -bracketing of F where ε ≤ 1 /n ≤ 1 / 2 ; we will study each element and conclude with a union bound. For any lower bracket l and upper bracket u in the bracketing (note l, u need not correspond to the same bracket). Recall that α l h +1 ( x, a ) := ∫ l h +1 ( z | x, a ) , so we have 1 -ε ≤ α l h +1 ≤ 1 since l is a lower ε -bracket of distributions. Therefore, we have

where ν h,i is the distribution of data from i -th round and time h . Note that ν h,i ( x, a, c, x ′ ) = d h,i ( x, a ) C h ( c | x, a ) P h ( x ′ | x, a ) for some distribution d h,i ( x, a ) . Now focus on each i , so for all i , we have

Therefore,

Thus, by Markov's inequality, w.p. at least 1 -δ , we have

To conclude, apply union bound to get this result for all brackets.

For the remainder of this section, we assume the policy class Π is finite. However, it is possible to extend our results using policy covers in the Hamming distance; in that case, log | Π | would be replaced by the log covering number or entropy integral of Π [as in Zhou et al., 2023, Kallus et al., 2022]. We note that for the online case, we rely on the assumption that for any f ∈ F we have π f ∈ Π , where recall that π f h ( x ) = arg min a ¯ f h ( x, a ) . This is because T ⋆,D is not a contraction so we cannot operate with T ⋆,D directly and instead operate with T π f ,D . We highlight that this assumption is automatically satisfied in tabular MDPs, since the whole policy space is finite, and log | Π | = O ( X log( A )) is lower order compared to log of the bracketing entropy of F tab , which is O ( X 2 A 2 ) . In contrast, in non-distributional methods such as GOLF, the regular Bellman optimality operator is a contraction so standard Lipschitz arguments for covering go through. We note that it is also possible to construct covers of F in the Hellinger distance, but the metric entropy of F tab 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

Then, construct the version space as

Theorem F.2. Fix any δ ∈ (0 , 1) and suppose Assumption 5.1. Set β = log( KH · N [] ( K -1 , F , ∥ · ∥ ∞ ) | Π | /δ ) . Then, w.p. at least 1 -δ , the following holds:

The optimal cost distribution is in the version space, i.e., Z ⋆ ∈ F k .

For all f ∈ F and h ∈ [ H ]

Proof. First, we want to verify that Z ⋆ ∈ F k . Let f ↓ be the lower bracket of Z ⋆ and set g = MLE f ↓ ,π ⋆ h ∈ F ; note π ⋆ = π Z ⋆ . By Theorem F.1, we have ∑ k i =1 log MLE f ↓ ,π ⋆ h ( z f ↓ ,π ⋆ h,i | x h,i , a h,i ) -log T π ⋆ ,D h f ↓ h +1 ( z f ↓ ,π ⋆ h,i | x h,i , a h,i ) ≤ O ( β ) . Therefore, noting that Z ⋆ h = T π ⋆ ,D h Z ⋆ h +1 ≥ T π ⋆ ,D h f ↓ h +1 shows that Z ⋆ h ∈ F f ↓ ,π ⋆ ,h for every h , implying that Z ⋆ ∈ F k .

For the second claim, fix any f ∈ F k and h ∈ [ H ] . Then,

The β comes from Theorem E.5, and for ε , we used the fact that H 2 ≤ H ≤ TV , and

since for any x, a , we have ∫ z ∣ ∣ ∣ ˜ f ↓ h +1 ( z | x, a ) -f h +1 ( z | x, a ) ∣ ∣ ∣ ≤ 3 ε . There are two cases. If ˜ f ↓ h +1 ( z | x, a ) ≥ f h +1 ( z | x, a ) , then ˜ f ↓ h +1 ( z | x, a ) -f h +1 ( z | x, a ) ≤ (1 -ε ) -1 f ↓ h +1 ( z | x, a ) -f h +1 ( z | x, a ) ≤ 2 εf h +1 ( z | x, a ) since (1 -ε ) -1 ≤ 1 + 2 ε . If ˜ f ↓ h +1 ( z | x, a ) < f h +1 ( z | x, a ) , then f h +1 ( z | x, a ) -˜ f ↓ h +1 ( z | x, a ) ≤ f h +1 ( z | x, a ) -f ↓ h +1 ( z | x, a ) ≤ ε . Thus, ∫ z max(2 εf h +1 ( z | x, a ) , ε ) ≤ ∫ z 2 εf h +1 ( z | x, a ) + ε = 3 ε . Thus, setting ε = 1 /K gives

For the offline setting, fix any π and define its general version space as,

Theorem F.3. Fix any δ ∈ (0 , 1) and suppose Assumption 5.1. Set β = log( H | Π | · N [] (( n |Y| ) -1 , F , ∥ · ∥ ∞ ) /δ ) . Then, w.p. at least 1 -δ , the following holds for all policies π ∈ Π :

The policy cost distribution is in the version space, i.e., Z π ∈ F π .

Any function in the version space has bounded triangular discrimination with the ground truth data-generating distribution, i.e., for all f ∈ F π and h ∈ [ H ] ,

Proof. The proof is the same as in Theorem F.2, but instead of π f , we fix any π .

G The ℓ p distributional eluder dimension

Let S denote any input space (for example, we will later instantiate S = X or S = X × A ). Let Ψ denote a set of functions mapping from S → R . Let D be a set of distributions on S .

Recall the definition of ε -independent sequence (of distributions) from Jin et al. [2021a].

Definition G.1 ( ℓ 2 -independent sequence) . Adistribution ν ∈ D is ( ε, ℓ 2 ) -independent of a sequence { d (1) , . . . , d ( n ) } ⊂ D if there exists ψ ∈ Ψ such that | E ν ψ | > ε and also √ ∑ n i =1 ( E d ( i ) ψ ) 2 ≤ ε .

Note that the definition is on sequences of distributions, which generalizes the original definition on sequences of points from Russo and Van Roy [2013].

We now generalize the above definition for the general ℓ p norm.

Definition G.2 ( ℓ p -independent sequence) . Adistribution ν ∈ D is ( ε, ℓ p ) -independent of a sequence { d (1) , . . . , d ( n ) } ⊂ D if there exists ψ ∈ Ψ such that | E ν ψ | > ε and also ∑ n i =1 | E d ( i ) ψ | p ≤ ε p .

Using the definition of independent sequences established so far, we define the ℓ p distributional eluder dimension.

Definition G.3 ( ℓ p -distributional eluder dimension) . For any p , define the ℓ p -distributional eluder dimension (denoted by DE p (Ψ , D , ε ) ) as the length of the longest sequence { d (1) , . . . , d ( d ) } ⊂ D such that there exists ε ′ ≥ ε , such that for all t ∈ [ d ] , d ( t ) is ( ε ′ , ℓ p ) -independent of d (1) , . . . , d ( t -1) .

Of particular interest to us is the ℓ 1 case. We show that the ℓ 1 eluder dimension is dominated by the ℓ 2 eluder dimension of Jin et al. [2021a].

Lemma 5.4. For any Ψ , D and ε > 0 , we have DE 1 (Ψ , D , ε ) ≤ DE 2 (Ψ , D , ε ) .

Proof. Since √∑ i x 2 i ≤ ∑ i | x i | , we have that any witness (long independent sequence) for ℓ 1 is also a witness for ℓ 2 . So, the maximum length of the ℓ 2 witnesses is longer than the ℓ 1 witnesses. Liu et 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 ℓ 1 distributional eluder dimension.

Theorem 5.3. Let C := sup d ∈D ,f ∈ Ψ | E d f | be the envelope. Fix any K ∈ N and sequences f (1) , . . . , f ( K ) ⊆ Ψ , d (1) , . . . , d ( K ) ⊆ D . Let β be a constant such that for all k ∈ [ K ] , we have, ∑ k -1 i =1 ∣ ∣ E d ( i ) f ( k ) ∣ ∣ ≤ β. Then, for all k ∈ [ K ] , we have

Proof. For any Γ ⊂ D , ν ∈ D , and 0 < ε ≤ 1 , let L ( ν, Γ , ε ) denote the number of disjoint subsets of Γ such that each subset is ε -dependent of ν , i.e. , for all such disjoint subsets of Γ , it is not the case that ν is ( ε, ℓ 1 ) -independent of each subset.

Fact 1: For any ε , if ∣ ∣ E d ( k ) f ( k ) ∣ ∣ > ε for some k ∈ [ K ] , then L ( d ( k ) , d (1: k -1) , ε ) < β/ε . By definition of L := L ( d ( k ) , d (1: k -1) , ε ) , there exist disjoint subsequences G (1) , . . . , G ( L ) of d (1: k -1) such that each subsequence G ( i ) satisfies ∑ d ∈ G ( i ) ∣ ∣ E d f ( k ) ∣ ∣ > ε . Therefore, summing over all subsequences, we have Lε < ∑ k -1 i =1 ∣ ∣ E d ( i ) f ( k ) ∣ ∣ ≤ β , where the β inequality comes from the premise. This proves Fact 1.

Fact 2: For any ε and any sequence { ν (1) , . . . , ν ( κ ) } ⊂ D , there exists j ∈ [ κ ] such that L ( ν ( j ) , ν (1: j -1) , ε ) ≥ J := ⌊ ( κ -1) / DE 1 (Ψ , D , ε ) ⌋ . If J = 0 , the claim is vacuously true. Otherwise, consider the following algorithm for finding the j :

Step 2) If ν ( j ) is ε -dependent on all of G ( i ) , i ∈ [ J ] , then the claim is proven and terminate.

Step 3) Otherwise, there exists some G ( i ) , i ∈ [ J ] such that ν ( j ) is ε -independent of it. Append ν ( j ) to G ( i ) , i.e. , G ( i ) = G ( i ) +[ ν ( j ) ] . Increment j = j +1 and go back to Step 2.

Hence, we need to argue this process terminates at Step 2 before j gets to κ +1 . We prove this by contradiction: assume j gets to κ +1 . Let i ∈ [ J ] be such that G ( i ) has the most elements (break ties arbitrarily). Since κ = ∑ J i =1 ∣ ∣ G ( i ) ∣ ∣ ≤ J ∣ ∣ G ( i ) ∣ ∣ , we have that ∣ ∣ G ( i ) ∣ ∣ ≥ κ/J ≥ κ κ -1 DE 1 (Ψ , D , ε ) > DE 1 (Ψ , D , ε ) , where we've also used the definition of J . By construction, G ( i ) is an ε -eluder sequence, i.e. , it is a sequence such that each element is ε -independent of its predecessors. However, this is a contradiction because its size is greater than DE 1 (Ψ , D , ε ) . Therefore, this process terminates at Step 2 for some j , which is the witness for proving Fact 2.

Fact 3: For any ε and k ∈ [ K ] , we have ∑ k t =1 I [∣ ∣ E d ( t ) f ( t ) ∣ ∣ > ε ] ≤ ( βε -1 +1 ) DE 1 (Ψ , D , ε )+1 . Fix any ε and k ∈ [ K ] . Let { d ( i 1 ) , . . . , d ( i κ ) } be all the elements of d (1: k ) such that E d ( t ) f ( t ) > ε for t = i 1 , . . . , i κ . By Fact 2, there exists j ∈ [ κ ] such that L ( d ( i j ) , d ( i 1: j -1 ) , ε ) ≥ ⌊ ( κ -1) / DE 1 (Ψ , D , ε ) ⌋ . By Fact 1, we have L ( d ( i j ) , d (1: i j ) , ε ) ≤ β/ε . Finally notice that L ( d ( i j ) , d ( i 1: j -1 ) , ε ) ≤ L ( d ( i j ) , d (1: i j ) , ε ) since adding more elements can only create more ε -dependent-ofν disjoint subsets. Thus, combining these inequalities, we have ⌊ ( κ -1) / DE 1 (Ψ , D , ε ) ⌋ < β/ε . This implies κ ≤ ( βε -1 + 1) DE 1 (Ψ , D , ε ) + 1 , which proves Fact 3.

Finishing the proof

Fix any k ∈ [ K ] and ω > 0 . We have

This completes the proof.

G.1 Bounding V-type ℓ 2 eluder dimension in low-rank MDPs

Theorem G.4 (Bound of ℓ 2 distributional eluder for low-rank MDPs) . Suppose the MDP is a low-rank MDP. Let Ψ ⊂ X → [0 , 1] be any class of functions mapping X to [0 , 1] . Suppose

Proof. If h = 1 , then D is a singleton. Hence, DE 2 (Ψ , D , ε ) ≤ 1 . Hence, suppose h ≥ 2 ; set h := h -1 and we will focus on d π h +1 in the remainder. Suppose { d ( k ) , f ( k ) } k ∈ [ T ] is any sequence such for all k ∈ [ T ] , we have that ( d ( k ) , f ( k ) ) is ( ε, ℓ 2 ) -independent of its predecessors. For any k , set Σ k = ∑ k -1 i =1 E d ( i ) [ ϕ ⋆ h ( x h , a h )] E d ( i ) [ ϕ ⋆ h ( x h , a h )] ⊤ + λI . Then, we have

Focusing on the second term,

Thus, we have shown that

Then, by the independent sequence assumption, we have

For a reference of the elliptical potential, see Uehara et al. [2021, Lemmas 19&20]. Rearranging, we have √ T < 2 √ d log(1 + T / ε 2 ) , which implies

By applying Lemma G.5, we have T ≤ 24 d log(1 + 4 d / ε 2 ) . This concludes the proof.

Lemma G.5. Let c 1 , c 2 ≥ 1 be constants. Let x ≥ 0 be a solution to x ≤ c 1 log(1 + c 2 x ) . Then, we necessarily have x ≤ 6 c 1 log(1 + c 1 c 2 ) .

Proof. Using change of variables B = x c 1 , we have the inequality is equivalent to B ≤ log(1 + B · c 1 c 2 ) . Take exp of both sides to get exp( B ) ≤ αB +1 where α = c 1 c 2 . From Step 3 of the proof of Russo and Van Roy [2013, Proposition 6], we have B ≤ e e -1 e e -1 (log(1 + α ) + log( e/ ( e -1))) ≤ 3(log(1 + c 1 c 2 ) + 1) . Hence, x ≤ c 1 · 3(log(1 + c 1 c 2 ) + 1) .

G.2 Bounding Q-type ℓ 2 eluder dimension in tabular MDPs

Theorem G.6 (Bound of ℓ 2 distributional eluder for tabular MDPs) . Suppose the MDP is a tabular MDP. Let Ψ ⊂ X × A → [0 , 1] be any class of functions mapping X × A to [0 , 1] . Suppose D be any set of distributions. Then, we have

Proof. Suppose { d ( k ) , f ( k ) } k ∈ [ T ] is any sequence such for all k ∈ [ T ] , we have that ( d ( k ) , f ( k ) ) is ( ε, ℓ 2 ) -independent of its predecessors. Since the MDP is tabular, we can interpret d ( k ) , f ( k ) as SA -dimensional vectors. For any k , set Σ k = ∑ k -1 i =1 d ( i ) ( d ( i ) ) ⊤ + λI . Then, we have

Focusing on the second term, we have

Thus, we have

Rearranging, we have √ T < 2 √ SA log(1 + T/ε 2 ) , which implies T ≤ 4 SA log(1 + T/ε 2 ) . Then by applying Lemma G.5, we have T ≤ 24 SA log(1 + 4 SA / ε 2 ) . This concludes the proof.

H Proofs for Online RL

H.1 Preliminary Lemmas

Lemma H.1. For any policy π , conditional distribution d and h ∈ [ H ] , we have

Lemma H.2 (Performance Difference Lemma (PDL)) . For any f : ( X × A → R ) H and policies π, π ′ , we have

Proof. We proceed by inducting on the following claim: for all h = H +1 , H, . . . , 1 ,

The base case of H +1 is trivially true as everything is 0 . Now fix any h and suppose the IH at h +1 is true. Then

By the IH, the first term is equal to ∑ H t = h +1 E π,x h [ T π ′ t f t +1 ( x t , a t ) -f t ( x t , π ′ ) ] . The second term is exactly E π,x h [ T π ′ h f h +1 ( x h , a h ) -f h ( x h , π ′ ) ] , which concludes the proof.

Proof.

H.2 Proof of Small-Loss Regret and PAC Bounds

Recall that we defined the function class and distribution class, for each h , as

Also, define the ' V -type' analogs as follows, which will be useful for PAC instead of regret bounds.

Let us also overload notation for the eluder dimensions as

Before we prove the following main theorem, a couple of remarks are in order:

Recall that by Theorem G.6, we have DE 1 ( ε ) ≤ O ( SA log( SA/ε )) and by Theorem G.4, we have DE 1 ,v ( ε ) ≤ O ( d log( d/ε )) . This shows that the Eluder dimension in terms in Theorem 5.5 are appropriately bounded.

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/ε ) + log | Φ | ) where Φ is a realizable class for ϕ ⋆ . This shows that the BC assumption of Theorem 5.5 is satisfied and β is appropriately bounded for low-rank MDPs.

Taken together, these two points imply that we have a small-loss PAC bound for low-rank MDPs: concretely, we have V ¯ π -V ⋆ ≤ ˜ O ( dH √ AV ⋆ log | Φ | K + d 2 H 2 A 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 ℓ ∞ bracketing number, i.e. , β = log( HKN [] (1 /K, F , ℓ ∞ ) /δ ) .

Theorem 5.5. Suppose DistBC holds (Assumption 5.1). For any δ ∈ (0 , 1) , w.p. at least 1 -δ , running O-DISCO with β = log( HK |F| /δ ) guarantees the following regret bound,

If UAE = TRUE (Algorithm 4), then the learned mixture policy ¯ π is guaranteed to satisfy,

Proof. For shorthand, let δ h,k ( x, a ) := D △ ( f ( k ) h ( x, a ) ∥ T ⋆,D h f ( k ) h +1 ( x, a )) and ∆ k := ∑ H h =1 E π k [ δ h,k ( x h , a h )] . Notice that since π k h +1 ( x ) = arg min a ¯ f ( k ) h +1 ( x, a ) , we have T π k ,D h f ( k ) h +1 ( x, a ) = T ⋆,D h f ( k ) h +1 ( x, a ) , so δ h,k ( x, a ) = D △ ( f ( k ) h ( x, a ) ∥ T π k ,D h f ( k ) h +1 ( x, a )) as well.

By Theorem F.2, we have the following two facts for all k ∈ [ K ] ,

(i) Optimism: min a ¯ f ( k ) 1 ( x 1 , a ) ≤ V ⋆ (since Z ⋆ ∈ F k ) and

(ii) Low training error: for all h , we have

If UAE=FALSE. ∑ i<k E π i [ δ h,k ( s h , a h )] ≤ 240 β .

The 240 comes from the constants of Theorem F.2 and the fact that D △ ( a, b ) ≤ 4 H 2 ( a, b ) for all distributions a, b .

In ⋆ , we used Cauchy Schwartz. Setting η = 4 H and rearranging, we have

√

Plugging this into ⋆ , and noting 104 e +17 ≤ 300 , we have

Thus, summing the instantaneous regrets over all episodes, we get

Last step: bounding ∑ k ∆ k . In this final step, we invoke the pigeonhole property of the eluder dimension, as proven in Theorem 5.3. Note that the precondition of 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 is at most 1 , and we will also pick ε = 1 /K .

On one hand, if UAE=FALSE, then,

On the other hand, if UAE=TRUE, then, we use the V-type analogs,

This concludes the proof for both the regret and PAC bounds.

Lemma H.3 (Self-bounding lemma) . Let f ∈ F and let π be any policy. Let us denote δ h ( x, a ) := D △ ( f h ( x, a ) ∥ T π,D h f h +1 ( x, a )) . Then, for all h ∈ [ H ] , for all x h , a h , we have

Proof. We prove the following refined subclaim inductively: for all h ∈ [ H ] , for all x h , a h , we have

For H +1 this is trivially true. Now fix any h and suppose IH is true for h +1 . By Eq. ( △ 2 ), for any h, x h , a h , we have,

In particular, we have that

which proves the inductive claim. Noting that ∑ H t =1 (1 + 1 /H ) t ≤ 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 δ ∈ (0 , 1) and set β = log( HK |F| /δ ) . Then, w.p. at least 1 -δ ,

In terms of H,X,A,K scaling, our bound matches that of GOLF [Xie et al., 2023] and is only a H factor looser than that of the minimax lower bound ˜ O ( √ XAK ) . The key benefit over prior bounds is that our leading term scales with the minimum cost of the problem V ⋆ . For example, if V ⋆ ≈ 0 , O-DISCO attains O (log K ) regret while uniform regret bounds are lower bounded by Ω( √ K ) . Compared to the minimax-optimal UCBVI [Azar et al., 2017], one weakness of our theorem is that it needs a F satisfying BC. Fortunately, in tabular MDPs where cost is only revealed at the last step from a known distribution, we can choose F tab as described in Wu et al. [2023, Lemma 4.15] to automatically satisfy BC. By extending our theory via bracketing entropy (Appendix F), we can derive that F tab yields β = O ( X 2 A 2 log( XAHK/δ )) . We note that if costs are unknown but discrete, it is possible to construct a BC function class with β scaling as O ( X 2 A 2 log( nXAHK/δ )) where n is the maximum number of possible cumulative costs.

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 ϕ ⋆ ( x, 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 π h ( z | x, 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 5.1. For any δ ∈ (0 , 1) , w.p. at least 1 -δ , running P-DISCO with β = log( H | Π ||F| /δ ) learns a policy ̂ π that enjoys the following PAC bound with respect to any comparator policy ˜ π ∈ Π :

Proof of Theorem 6.1. For shorthand, let δ π h ( x, a ) = D △ ( f π h ( x, a ) ∥ T π,D h f π h +1 ( x, a )) and ∆ π = ∑ H h =1 E π [ δ π h ( x h , a h )] . Also, let f ( x, π ) = E a ∼ π ( x ) [ f ( x, a )] .

By Theorem F.3, we have the following two facts, for all π ∈ Π , (i) Pessimism: V π ≤ ¯ f π 1 ( x 1 , π ) (since Z π ∈ F π ) for all π ∈ Π , and

(ii) E ν h [ δ π h ( x h , a h )] ≤ β ′ N -1 for all h where Theorem F.3 and the fact that D △ ≤ 4 H 2 certifies that β ′ = 240 β is sufficient.

With these two facts, we can bound the suboptimality of ̂ π as follows:

Finally, we can bound ∆ ˜ π by a change of measure,

Therefore,

J Extension: Small-Return Bounds

In this section, we show that O-DISCO and P-DISCO can also be used to obtain small-return bounds. Compared to the algorithms presented in the main text for minimizing cost, we simply have to replace min with max (and vice versa) for maximizing reward, i.e. , see Appendix B 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 δ ∈ (0 , 1) and set β = log( HK |F| /δ ) and β ′ = 60 β . Then, w.p. at least 1 -δ , running O-DISCO (Algorithm 4) with UAE = FALSE yields the following small-loss regret bound,

If instead UAE = TRUE , the outputted policy ¯ π enjoys the following small-loss PAC bound,

Proof. Adopt the same notation as in the proof of Theorem 5.5. By Theorem F.2, we have the following two facts for all k ∈ [ K ] ,

(i) Optimism: V ⋆ ≤ max a ¯ f ( k ) 1 ( x 1 , a ) (since Z ⋆ ∈ F k ) and (ii) ∑ i<k E π i [ δ h,k ( s h , a h )] ≤ β ′ for all h . If UAE=TRUE, then a h is sampled from unif( A ) rather than π i , i.e. , we have ∑ i<k E s h ∼ π i ,a h ∼ unif( A ) [ δ h,k ( s h , a h )] ≤ β ′ , where β ′ ≲ β . Theorem F.2 certifies that β ′ = 60 β is sufficient.

Fix any episode k ∈ [ K ] . Then,

Thus, summing the instantaneous regrets over all episodes, we get

The bounds for ∆ k are the same as in Theorem 5.5.

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-Schwartz step to rearrange terms at the step labelled ♣ . However, in the reward-maximizing setting, we simply bound V π k ≤ 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 δ ∈ (0 , 1) and set β = log( H | Π ||F| /δ ) . Then, w.p. at least 1 -δ , P-DISCO (Algorithm 4) learns a policy ̂ π such that for any comparator policy π ∈ Π , we have

Proof of Theorem J.2. Adopt the same notation as in the proof of Theorem 6.1. By Theorem F.3, we have the following two facts, for all π ∈ Π ,

(i) Pessimism: ¯ f π 1 ( x 1 , π ) ≤ V π (since Z π ∈ F π ) for all π ∈ Π , and

(ii)

E

ν

h

[

δ

π

h

(

x

h

, a

h

)]

β

N

for all

h

where

β

60

β

.

With these two facts, we can bound the suboptimality of ̂ π as follows:

From here, the same argument in the proof of Theorem 6.1 finishes the proof.

K Experiment Details

Experiment Settings

In our experiments, as outlined in Foster and Krishnamurthy [2021], our γ learning rate at each time step t is set to γ t = γ 0 t p where γ 0 and p are hyperparameters. We use batch sizes of 32 samples per episode, and the King County and Prudential experiments run 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 as 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 use the WandB (Weights and Biases) library to run 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 a simple 2 hidden-layer neural network for the Prudential Life Insurance and King County Housing datasets. For DISTCB, the oracle's output layer has size AC where A is the number of actions and C is 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 to represent cost probabilities. For SquareCB and FastCB, the output size is A because there is just a single prediction associated with each action. As per Foster and Krishnamurthy [2021], a sigmoid function is applied to this output layer. All experiments were implemented using 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.

′

-

1

′

˜

Prudential Life Insurance This dataset is from the Prudential Life Insurance Kaggle competition [Montoya et al., 2015]. It is 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 not getting a sale. Otherwise, the model's prediction is charging too little for the customer. To reiterate, the cost in this case is . 1 ∗ ( y -ˆ y ) where y is the actual risk level, and ˆ y is the predicted risk level.

King County Housing The King County housing dataset is also used in 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, as it lets us delineate 3 possible costs that DISTCB can learn. Without the super class, the cost construction would be a pure binary of 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 a . 086 advantage over the last 100 episodes and a . 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\n1 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\nQUESTIONS:\nsee weakness.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nThe authors do not describe any specific limitation of this paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nNone.\n\nLIMITATIONS:\nThe paper includes a brief discussion of the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nSignificance 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.\n\nWEAKNESSES:\nClarity:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-12"}