SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Kaiwen Wang Kevin Zhou Runzhe Wu Nathan Kallus Wen Sun
{kw437,klz23,rw646,kallus,ws455}@cornell.edu
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 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]. For example in linear MDPs, typical worst-case regret bounds scale on the order of  $\mathrm{poly}(d,H)\sqrt{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  $\mathrm{poly}(d,H)\sqrt{K\cdot V^{\star}} +\mathrm{poly}(d,H)\log (K)$ , where  $V^{\star} = \min_{\pi}V^{\pi}$  is the optimal expected cumulative cost for the problem. We assume cumulative costs are normalized in [0, 1] without loss of generality. As  $V^{\star}$  becomes negligible (approaches 0), the first term vanishes and the small-loss bound yields a faster convergence rate of  $\mathcal{O}(\mathrm{poly}(d,H)\log (K))$ , compared to the  $\mathcal{O}(\mathrm{poly}(d,H)\sqrt{K})$  rate in standard uniform bounds. Since we always have  $V^{\star}\leq 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  $\ell_1$  distributional eluder dimension, which generalizes the  $\ell_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 depend on the former. Finally, we find that our distributional CB algorithm empirically outperforms existing approaches in three challenging CB tasks.
1. As warm-up, we propose a distributional CB algorithm and prove that it obtains a small-loss regret bound (Section 4). We empirically demonstrate it outperforms state-of-the-art CB algorithms in three challenging benchmark tasks (Section 7).
2. We propose a distributional online RL algorithm that enjoys small-loss bounds in settings with low  $\ell_1$  distributional eluder dimension, which we show can always capture low-rank MDPs. The  $\ell_1$  distributional eluder dimension may be of independent interest (Section 5).
3. We propose a distributional offline RL algorithm and prove that it obtains the first small-loss bounds in the offline setting. Our small-loss guarantee exhibits a novel robustness to bad coverage, which implies strong improvement over more policies than existing results in the literature (Section 6).
In sum, we show that DistRL can yield small-loss bounds in both online and offline RL, which provide a concrete theoretical justification for the benefits of distribution learning in decision making.
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" generally refers to bounds that scale with the optimal value, either the maximum reward or the minimum cost. To highlight that we are minimizing cost, we call our bounds "small-loss".
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.)  $\delta$  contributes at most  $R^{\star}\delta$  to the regret $^{2}$ . When minimizing costs, receiving the worst loss 1 w.p.  $\delta$  may induce large regret relative to  $L^{\star}$  if  $L^{\star}$  is small. To illustrate, if  $R^{\star} = 0$  then all policies are optimal, so no learning is needed and the small-return bound is vacuous. Yet if  $L^{\star} = 0$ , sub-optimal policies may have a large gap from  $L^{\star}$ , 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^{\star}$  in the reward-maximizing setting, i.e.,  $\mathrm{Var}(\sum_{h}r_{h}\mid \pi^{k})\leq c\cdot V^{\pi^{k}}\leq c\cdot V^{\star}$ . 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.
As warmup, we begin with the contextual bandit problem with an arbitrary context space  $\mathcal{X}$ , finite action space  $\mathcal{A}$  with size  $A$  and conditional cost distributions  $C: \mathcal{X} \times \mathcal{A} \to \Delta([0,1])$ . Throughout, we fix some dominating measure  $\lambda$  on  $[0,1]$  (e.g., Lebesgue for continuous or counting for discrete) and let  $\Delta([0,1])$  be all distributions on  $[0,1]$  that are absolutely continuous with respect to  $\lambda$ . We identify such a distribution with its density with respect to  $\lambda$ , and we also write  $C(y \mid x, a)$  for  $(C(x, a))(y)$ . Let  $K$  denote the number of episodes. At each episode  $k \in [K]$ , the learner observes a context  $x_k \in \mathcal{X}$ , samples an action  $a_k \in \mathcal{A}$ , and then receives a cost  $c_t \sim C(x_t, a_t)$ , which we assume to be normalized, i.e.,  $c_t \in [0,1]$ . The goal is to design a learner that attains low regret with high probability, where regret is defined as
$$
\operatorname {R e g r e t} _ {\mathrm {C B}} (K) = \sum_ {k = 1} ^ {K} \bar {C} \left(x _ {k}, a _ {k}\right) - \bar {C} \left(x _ {k}, \pi^ {\star} \left(x _ {k}\right)\right),
$$
where  $\bar{f} = \int yf(y)\mathrm{d}\lambda (y)$  for any  $f\in \Delta ([0,1])$  and  $\pi^{\star}(x_k) = \arg \min_{a\in A}\bar{C} (x_k,a)$ .
The focus of this paper is reinforcement learning (RL) under the Markov Decision Process (MDP) model, with observation space  $\mathcal{X}$ , finite action space  $\mathcal{A}$  with size  $A$ , horizon  $H$ , transition kernels  $P_{h}:\mathcal{X}\times \mathcal{A}\to \Delta (\mathcal{X})$  and cost distributions  $C_h:\mathcal{X}\times \mathcal{A}\rightarrow \Delta ([0,1])$  at each step  $h\in [H]$ . We start with the Online RL setting, which proceeds over  $K$  episodes as follows: at each episode  $k\in [K]$ , the learner plays a policy  $\pi^k\in [\mathcal{X}\to \Delta (\mathcal{A})]^H$ ; we start from a fixed initial state  $x_{1}$ ; then for each  $h = 1,2,\ldots ,H$ , the policy samples an action  $a_{h}\sim \pi_{h}^{k}(x_{h})$ , receives a cost  $c_{h}\sim C_{h}(x_{h},a_{h})$ , and transitions to the next state  $x_{h + 1}\sim P_h(x_h,a_h)$ . Our goal is to compete with the optimal policy that minimizes expected the loss, i.e.,  $\pi^{\star}\in \arg \min_{\pi \in \Pi}V^{\pi}$  where  $V^{\pi} = \mathbb{E}_{\pi}\Big[\sum_{h = 1}^{H}c_{h}\Big]$ . Regret bounds aim to control the learner's regret with high probability, where regret is defined as,
$$
\operatorname {R e g r e t} _ {\mathrm {R L}} (K) = \sum_ {k = 1} ^ {K} V ^ {\pi^ {k}} - V ^ {\star}.
$$
If the algorithm returns a single policy  $\widehat{\pi}$ , it is desirable to obtain a Probably Approximately Correct (PAC) bound on the sub-optimality of  $\widehat{\pi}$ , i.e.,  $V^{\widehat{\pi}} - V^{\star}$ .
2Assume rewards/losses in [0, 1] and  $R^{\star} / L^{\star}$  is the maximum/minimum expected reward/loss.
The third setting we study is Offline RL, where instead of needing to actively explore and collect data ourselves, we are given  $H$  datasets  $\mathcal{D}_1,\mathcal{D}_2,\ldots ,\mathcal{D}_H$  to learn a good policy  $\widehat{\pi}$ . Each  $\mathcal{D}_h$  contains  $N$  i.i.d. samples  $(x_{h,i},a_{h,i},c_{h,i},x_{h,i}^{\prime})$  from the process  $(x_{h,i},a_{h,i})\sim \nu_h,c_{h,i}\sim C_h(x_{h,i},a_{h,i}),x_{h,i}^{\prime}\sim P_h(x_{h,i},a_{h,i})$ , where  $\nu_{h}\in \Delta (\mathcal{X}\times \mathcal{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  $\widehat{\pi}$  which should improve over the data generating process.
Distributional RL For a policy  $\pi$  and  $h\in [H]$ , let  $Z_h^\pi (x_h,a_h)\in \Delta ([0,1])$  denote the distribution of the loss-to-go  $\sum_{t = h}^{H}c_{t}$  conditioned on rolling in  $\pi$  from  $x_{h},a_{h}$ . The expectation of the above is  $Q_{h}^{\pi}(x_{h},a_{h}) = \bar{Z}_{h}^{\pi}(x_{h},a_{h})$  and  $V_{h}^{\pi}(x_{h}) = \mathbb{E}_{a_{h}\sim \pi_{h}(x_{h})}[Q_{h}^{\pi}(x_{h},a_{h})]$ . We use  $Z_h^\star ,Q_h^\star ,V_h^\star$  to denote these quantities with  $\pi^{\star}$ . Recall the regular Bellman operator acts on a function  $f:\mathcal{X}\times \mathcal{A}\to [0,1]$  as follows:  $\mathcal{T}_h^\pi f(x,a) = \bar{C}_h(x,a) + \mathbb{E}_{x'\sim P_h(x,a),a'\sim \pi (x')}[f(x',a')]$ . Analogously, the distributional Bellman operator [Morimura et al., 2012, Bellemare et al., 2017] acts on a conditional distribution  $d:\mathcal{X}\times \mathcal{A}\rightarrow \Delta ([0,1])$  as follows:  $\mathcal{T}_h^{\pi ,D}d(x,a)\stackrel {D}{=}C_h(x,a) + d(x',a')$ , where  $x^{\prime}\sim P_{h}(x,a),a^{\prime}\sim \pi (x^{\prime})$  and  $\underline{D}$  denotes equality of distributions. Another way to think about the distributional Bellman operator is that a sample  $z\sim \mathcal{T}_h^{\pi ,D}d(x,a)$  is generated as follow:  $z\coloneqq c + y$ , where  $c\sim C_h(x,a),x'\sim P_h(x,a),a'\sim \pi (x'),y\sim d(x',a')$ . We will also use the Bellman optimality operator  $\mathcal{T}_h^\star$  and its distributional variant  $\mathcal{T}_h^{\star ,D}$ , defined as follows:  $\mathcal{T}_h^\star f(x,a) = \bar{C}_h(x,a) + \mathbb{E}_{x'\sim P_h(x,a)}[\min_{a\in \mathcal{A}}f(x',a')]$  and  $\mathcal{T}_h^{\star ,D}d(x,a)\stackrel {D}{=}C_h(x,a) + d(x',a')$  where  $x^{\prime}\sim P_{h}(x,a),a^{\prime} = \arg \min_{a}\bar{d} (x^{\prime},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  $\widehat{f} \in \mathbb{R}^A$  and a parameter  $\gamma \in \mathbb{R}_{++}$ : setting  $b = \arg \min_{a \in A} \widehat{f}(a)$  as the best action with respect to the predictions, the weight for any other action  $a \neq b$  is,
$$
\operatorname {R e I G W} _ {\gamma} (\widehat {f}, \gamma) [ a ] := \frac {\widehat {f} (b)}{A \widehat {f} (b) + \gamma (\widehat {f} (a) - \widehat {f} (b))}, \tag {1}
$$
and the rest of the weight is allocated to  $b$ :  $\mathrm{ReIGW}_{\gamma}(\widehat{f},\gamma)[b] = 1 - \sum_{a\neq b}\mathrm{ReIGW}_{\gamma}(\widehat{f},\gamma)[a]$ .
Algorithm 1 Distributional CB (DISTCB)
1: Input: number of episodes  $K$ , failure probability  $\delta$ , ReIGW learning rate  $\gamma$ .
2: Initialize any cost distribution  $f^{(1)}$ .
3: for episode  $k = 1,2,\ldots ,K$  do
4: Observe context  $x_{k}$ .
5: Sample action  $a_{k}\sim p_{k} = \mathrm{ReIGW}(\bar{f}^{(k)}(x_{k},\cdot),\gamma)$  from Eq. (1).
6: Observe cost  $c_k \sim C(x_k, a_k)$  and update online MLE oracle with  $((x_k, a_k), c_k)$ .
7: end for
We propose Distributional Contextual Bandit (DISTCB) in Algorithm 1, a two-step procedure for each episode  $k \in [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.,  $\widehat{f}(a) = \bar{f}^{(k)}(x_{k}, a), \forall a \in \mathcal{A}$  (Line 5). Then, DISTCB updates  $f^{(k)}(\cdot \mid x_{k}, a_{k})$  by maximizing the log-likelihood to estimate the conditional cost distribution  $C(\cdot \mid x_{k}, a_{k})$  (Line 6). Formally, this second step is achieved via an online MLE oracle with a realizable distribution class  $\mathcal{F}_{CB} \subset \mathcal{X} \times \mathcal{A} \to \Delta([0,1])$ ; let  $\mathrm{Regret}_{\log}(K)$  be some upper bound on the log-likelihood regret for all possibly adaptive sequences  $\{x_{k}, a_{k}, c_{k}\}_{k \in [K]}$ ,
$$
\sum_ {k = 1} ^ {K} \log C \left(c _ {k} \mid x _ {k}, a _ {k}\right) - \log f ^ {(k)} \left(c _ {k} \mid x _ {k}, a _ {k}\right) \leq \operatorname {R e g r e t} _ {\log} (K).
$$
Under realizability,  $C \in \mathcal{F}_{CB}$ , we expect  $\mathrm{Regret}_{\log}(K) \in \mathcal{O}(\log(K))$ . For instance, if  $\mathcal{F}_{CB}$  is finite, exponentially weighted average forecaster guarantees  $\mathrm{Regret}_{\log}(K) \leq \log |\mathcal{F}_{CB}|$  [Cesa-Bianchi and Lugosi, 2006, Chapter 9]. We now state our main result for DISTCB.
Theorem 4.1. For any  $\delta \in (0,1)$ , w.p. at least  $1 - \delta$ , running DISTCB with  $\gamma = 10A \vee \sqrt{\frac{40A(C^{\star} + \log(1 / \delta))}{112\left(\mathrm{Regret}_{\log}(K) + \log(1 / \delta)\right)}}$  has regret scaling with  $C^{\star} = \sum_{k=1}^{K} \min_{a \in A} \bar{C}(x_k, a)$ .
$$
\operatorname {R e g r e t} _ {\mathrm {D i s t C B}} (K) \leq 2 3 2 \sqrt {A C ^ {\star} \operatorname {R e g r e t} _ {\log} (K) \log (1 / \delta)} + 2 3 0 0 A \left(\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta)\right).
$$
The dominant term scales with the optimal sum of costs  $\sqrt{C^{\star}}$  which shows that DISTCB obtains small-loss regret. DISTCB is also computationally efficient since each episode simply requires computing the ReIGW. FastCB is the only other computationally efficient CB algorithm with small-loss regret [Foster and Krishnamurthy, 2021, Theorem 1]. Our bound matches that of FastCB in terms of dependence on  $A$ ,  $C^{\star}$  and  $\log(1/\delta)$ . Our key difference with FastCB is the online supervised learning oracle: in DISTCB, we aim to learn the conditional cost distribution by maximizing log-likelihood, 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,
$$
\operatorname{Regret}_{\operatorname {DistCB}}(K)\lesssim \sum_{k = 1}^{K}\mathbb{E}_{a_{k}\sim p_{k}}\left[\frac{A}{\gamma}\bar{C} (s_{k},a_{k}) + \gamma \underbrace{\left(\bar{f}^{(k)}(s_{k},a_{k}) - \bar{C} (s_{k},a_{k})\right)^{2}}_{\star}\right].
$$
For any distributions  $f, g \in \Delta([0,1])$ , their triangular discrimination $^3$  is defined as  $D_{\triangle}(f \parallel g) := \int \frac{(f(y) - g(y))^2}{f(y) + g(y)} \mathrm{d}\lambda(y)$ . The key insight is that  $\star$  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 \leq y$  for  $y \in [0,1]$ , we have  $\bar{f} - \bar{g} = \int y(f(y) - g(y)) \mathrm{d}\lambda(y) \leq \sqrt{\int y(f(y) + g(y)) \mathrm{d}\lambda(y)} \sqrt{\int \frac{(f(y) - g(y))^2}{f(y) + g(y)} \mathrm{d}\lambda(y)}$ , and hence,
$$
\left| \bar {f} - \bar {g} \right| \leq \sqrt {\left(\bar {f} + \bar {g}\right) D _ {\triangle} (f \| g)}. \tag {△1}
$$
So, Eq.  $(\triangle_1)$  implies that  $\star$  is bounded by  $D_{\triangle}(f^{(k)}(s_k,a_k)\parallel C(s_k,a_k))$ . Since  $D_{\triangle}$  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 - \delta$ , we have
$$
\operatorname {R e g r e t} _ {\operatorname {D i s t C B}} (K) \lesssim \sum_ {k = 1} ^ {K} \frac {A}{\gamma} \big (\bar {C} (s _ {k}, a _ {k}) + \log (1 / \delta) \big) + \gamma \big (\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta) \big).
$$
From here, we just need to rearrange terms and set the correct  $\gamma$ . 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  $\mathcal{F} \subseteq (\mathcal{X} \times \mathcal{A} \to \Delta([0,1]))^H$  where each element  $f \in \mathcal{F}$  is a tuple  $f = (f_1, \ldots, f_H)$  such that each  $f_h$  is a candidate estimator for  $Z_h^\star$ , the distribution of loss-to-go  $\sum_{t=h}^{H} c_t$  under  $\pi^\star$ . For notation,  $f_{H+1}(x,a) = \delta_0$  denotes the dirac at zero for all  $x, a$ .
We now present our Optimistic Distributional Confidence set Optimization (O-DISCO) algorithm in Algorithm 2, consisting of three key steps per episode. At episode  $k \in [K]$ , O-DISCO first identifies the  $f^{(k)}$  with the minimal expected value at  $h = 1$  over the previous confidence set  $\mathcal{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  $\pi^k$  with respect to the mean of  $f^{(k)}$  (Line 6). Finally, O-DISCO constructs a
Triangular discrimination is also known as Vincze-Le Cam divergence [Vincze, 1981, Le Cam, 2012].
Algorithm 2 Optimistic Distributional Confidence set Optimization (O-DISCO)