SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1: Input: number of episodes  $K$ , distribution class  $\mathcal{F}$ , threshold  $\beta$ .
2: Initialize  $\mathcal{D}_{h,0} \gets \emptyset$  for all  $h \in [H]$ , and set  $\mathcal{F}_0 = \mathcal{F}$ .
3: for episode  $k = 1,2,\ldots ,K$  do
4: Set optimistic estimate  $f^{(k)} = \arg \min_{f\in \mathcal{F}_{k - 1}}\min_{a}\bar{f}_1(x_1,a)$ .
5: Set  $\pi_h^k (x) = \arg \min_a\bar{f}_h^{(k)}(x,a)$
6: Roll out  $\pi^k$  and obtain a trajectory  $x_{1,k},a_{1,k},c_{1,k},\ldots ,x_{H,k},a_{H,k},c_{H,k}$  For each  $h\in [H]$  , augment the dataset  $\mathcal{D}_{h,k} = \mathcal{D}_{h,k - 1}\cup \{(x_{h,k},a_{h,k},c_{h,k},x_{h + 1,k})\}$
7: For all  $(h,f)\in [H]\times \mathcal{F}$ , sample  $y_{h,i}^{f}\sim f_{h + 1}(x_{h,i}^{\prime},a^{\prime})$  and  $a^\prime = \arg \min_a\bar{f}_{h + 1}(x_{h,i}^\prime ,a)$  where  $(x_{h,i},a_{h,i},c_{h,i},x_{h,i}^{\prime})$  is the  $i$ -th datapoint of  $\mathcal{D}_{h,k}$ . Then, set  $z_{h,i}^{f} = c_{h,i} + y_{h,i}^{f}$  and define the confidence set
$$
\mathcal {F} _ {k} = \left\{f \in \mathcal {F}: \sum_ {i = 1} ^ {k} \log f _ {h} \left(z _ {h, i} ^ {f} \mid x _ {h, i}, a _ {h, i}\right) \geq \max  _ {g \in \mathcal {F} _ {h}} \sum_ {i = 1} ^ {k} \log g \left(z _ {h, i} ^ {f} \mid x _ {h, i}, a _ {h, i}\right) - 7 \beta , \forall h \in [ H ] \right\}.
$$
8: end for
9: Output:  $\bar{\pi} = \mathrm{unif}(\pi^{1:K})$ .
confidence set  $\mathcal{F}_k$  by including a function  $f$  if it exceeds a threshold on the log-likelihood objective using data  $z_{h,i}^{f}\sim \mathcal{T}_{h}^{\star ,D}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\in \mathcal{F}_k$  has the property that  $f_{h}$  is close-in-distribution to  $\mathcal{T}_h^{\star ,D}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  $\bar{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  $\pi, h \in [H]$ ,  $f_{h+1} \in \mathcal{F}_{h+1} \implies \mathcal{T}_h^{\pi,D} f_{h+1} \in \mathcal{F}_h$ .
5.1 The  $\ell_1$  Distributional Eluder Dimension
We now introduce the  $\ell_1$  distributional eluder dimension. Let  $S$  be an abstract input space, let  $\Psi$  be a set of functions mapping  $S\to \mathbb{R}$  and let  $\mathcal{D}$  be a set of distributions on  $S$ .
Definition 5.2 ( $\ell_p$ -distributional eluder dimension). For any function class  $\Psi \subseteq S \to \mathbb{R}$ , distribution class  $\mathcal{D} \subseteq \Delta(S)$  and  $\varepsilon > 0$ , the  $\ell_p$ -distributional eluder dimension (denoted by  $\mathrm{DE}_p(\Psi, \mathcal{D}, \varepsilon)$ ) is the length  $L$  of the longest sequence  $d^{(1)}, d^{(2)}, \ldots, d^{(L)} \subseteq \mathcal{D}$  such that there exists  $\varepsilon' \geq \varepsilon$ , such that for all  $t \in [L]$ , we have that there exists  $f \in \Psi$  such that  $|\mathbb{E}_{d^{(t)}}f| > \varepsilon$  and also  $\sum_{i=1}^{t-1} |\mathbb{E}_{d^{(i)}}f|^p \leq \varepsilon^p$ .
When  $p = 2$ , this is exactly the  $\ell_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  $\ell_1$  distributional eluder dimension.
Theorem 5.3. Let  $C \coloneqq \sup_{d \in \mathcal{D}, f \in \Psi} |\mathbb{E}_df|$  be the envelope. Fix any  $K \in \mathbb{N}$  and sequences  $f^{(1)}, \ldots, f^{(K)} \subseteq \Psi$ ,  $d^{(1)}, \ldots, d^{(K)} \subseteq \mathcal{D}$ . Let  $\beta$  be a constant such that for all  $k \in [K]$ , we have,  $\sum_{i=1}^{k-1} |\mathbb{E}_{d^{(i)}}f^{(k)}| \leq \beta$ . Then, for all  $k \in [K]$ , we have
$$
\sum_ {t = 1} ^ {k} \left| \mathbb {E} _ {d ^ {(t)}} f ^ {(t)} \right| \leq \inf  _ {0 <   \varepsilon \leq 1} \left\{\mathrm {D E} _ {1} (\Psi , \mathcal {D}, \varepsilon) (2 C + \beta \log (C / \varepsilon)) + k \varepsilon \right\}.
$$
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  $\ell_1$  dimension generalizes the original  $\ell_2$  dimension of Jin et al. [2021a].
Lemma 5.4. For any  $\Psi, \mathcal{D}$  and  $\varepsilon > 0$ , we have  $\mathrm{DE}_1(\Psi, \mathcal{D}, \varepsilon) \leq \mathrm{DE}_2(\Psi, \mathcal{D}, \varepsilon)$ .
Finally, we note that our distributional eluder dimension generalize the regular  $\ell_1$  eluder from Liu et al. [2022], which can be seen by taking  $\mathcal{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  $\mathcal{S} = \mathcal{X}\times \mathcal{A}$ . While low-rank MDPs are not captured by the "Q-type" dimension, they are captured by the "V-type" dimension where  $\mathcal{S} = \mathcal{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  $\pi^k$  for a single trajectory, partially roll-out  $\pi^k$  for  $H$  times where for each  $h\in [H]$ , we collect  $x_{h,k}\sim d_h^{\pi^k}$ , take a random action  $a_{h,k}\sim \mathrm{unif}(\mathcal{A})$ , observe  $c_{h,k}\sim C_h(x_{h,k},a_{h,k}),x_{h,k}^\prime \sim P_h(x_{h,k},a_{h,k})$  and augment the dataset  $\mathcal{D}_{h,k} = \mathcal{D}_{h,k - 1}\cup \{(x_{h,k},a_{h,k},c_{h,k},x_{h,k}^\prime)\}$ . 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  $\mathcal{D}_h = \{(x,a)\mapsto d_h^\pi (x,a):\pi \in \Pi \}$  and  $\Psi_{h} = \{(x,a)\mapsto D_{\triangle}(f(x,a)\parallel \mathcal{T}^{\star ,D}f(x,a)):f\in \mathcal{F}\}$ . Similarly, the V-type classes are  $\mathcal{D}_{h,v} = \{x\mapsto d_h^\pi (x):\pi \in \Pi \}$  and  $\Phi_{h,v} = \{x\mapsto \mathbb{E}_{a\sim \mathrm{Unif}(\mathcal{A})}[D_{\triangle}(f(x,a)\parallel \mathcal{T}^{\star ,D}f(x,a))]:f\in \mathcal{F}\}$ . Finally, define  $\mathrm{DE}_1(\varepsilon) = \max_h\mathrm{DE}_1(\Psi_h,\mathcal{D}_h,\varepsilon)$  and  $\mathrm{DE}_{1,v}(\varepsilon) = \max_h\mathrm{DE}_1(\Psi_{h,v},\mathcal{D}_{h,v},\varepsilon)$ .
Theorem 5.5. Suppose  $\text{DistBC}$  holds (Assumption 5.1). For any  $\delta \in (0,1)$ , w.p. at least  $1 - \delta$ , running O-DISCO with  $\beta = \log(HK|\mathcal{F}| / \delta)$  guarantees the following regret bound,
$$
\mathrm {R e g r e t} _ {\mathrm {O - D I S C O}} (K) \leq 1 6 0 H \sqrt {K V ^ {\star} \mathrm {D E} _ {1} (1 / K) \log (K) \beta} + 1 8 0 0 0 H ^ {2} \mathrm {D E} _ {1} (1 / K) \log (K) \beta .
$$
If UAE = TRUE (Algorithm 4), then the learned mixture policy  $\bar{\pi}$  is guaranteed to satisfy,
$$
V ^ {\bar {\pi}} - V ^ {\star} \leq 1 6 0 H \sqrt {\frac {A V ^ {\star} \mathrm {D E} _ {1 , v} (1 / K) \log (K) \beta}{K}} + \frac {1 8 0 0 0 H ^ {2} A \mathrm {D E} _ {1 , v} (1 / K) \log (K) \beta}{K}.
$$
Compared to prior bounds for GOLF [Jin et al., 2021a], the leading  $\sqrt{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^{\star}$ , giving our bounds the small-loss property. For example, if  $V^{\star} \leq \mathcal{O}(1 / \sqrt{K})$ , then our regret bound converges at a fast  $\mathcal{O}(H^{2} \mathrm{DE}_{1}(1 / K) \log (K)\beta)$  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<sup>4</sup>. 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 Fitted-  $Q$ -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  $\mathcal{G}_h$  such that  $f_{h + 1}\in \mathcal{F}_{h + 1}\Rightarrow \mathcal{T}_h^{\pi ,D}f_{h + 1}\in \mathcal{G}_h$  [as in Jin et al., 2021a, Assumption 14]. Simply replace  $\max_{g\in \mathcal{F}_h}$  in the confidence set construction with  $\max_{g\in \mathcal{G}_h}$ . While adding functions to  $\mathcal{F}$  may break BC (as BC is not monotonic), we can always augment  $\mathcal{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,  $\varepsilon$ -greedy suffices in problems that don't require deep and strategic exploration, i.e., a large myopic exploration gap [Dann et al., 2022]. With  $\varepsilon$ -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.
<sup>4</sup>In Appendix J, we show a slight modification of our approach also yields "small-return" bounds.
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}:\mathcal{X}\times \mathcal{A}\to \Delta (\mathcal{X})$  has rank  $d$  if there exist unknown features  $\phi_h^\star :\mathcal{X}\times \mathcal{A}\to \mathbb{R}^d,\mu_h^\star :\mathcal{X}\to \mathbb{R}^d$  such that  $P_{h}(x^{\prime}\mid x,a) = \phi_{h}^{\star}(x,a)^{\top}\mu_{h}^{\star}(x^{\prime})$  for all  $x,a,x'$ . Also, assume  $\max_{x,a}\| \phi_h^\star (x,a)\| _2\leq 1$  and  $\| \int g\mathrm{d}\mu_h^\star \| _2\leq \| g\|_\infty \sqrt{d}$  for all functions  $g:\mathcal{X}\rightarrow \mathbb{R}$ . The MDP is called low-rank if  $P_{h}$  is low-rank for all  $h\in [H]$ .
We now specialize Theorem 5.5 to low-rank MDPs with three key steps. First, we bound the V-type eluder dimension by  $\mathrm{DE}_{1,v}(\varepsilon) \leq \mathcal{O}(d\log (d / \varepsilon))$ , which is a known result that we reproduce in Theorem G.4. The next step requires access to a realizable  $\Phi$  class, i.e., for all  $h \in [H]$ ,  $\phi_h^\star \in \Phi$ , which is a standard assumption for low-rank MDPs [Agarwal et al., 2020, Uehara et al., 2021, Mhammedi et al., 2023]. Given the realizable  $\Phi$ , we can construct a specialized  $\mathcal{F}$  for the low-rank MDP:  $\mathcal{F}^{\mathrm{lin}} = \mathcal{F}_1^{\mathrm{lin}} \times \dots \times \mathcal{F}_H^{\mathrm{lin}} \times \mathcal{F}_{H + 1}^{\mathrm{lin}}$  where  $\mathcal{F}_{H + 1}^{\mathrm{lin}} = \{\delta_0\}$  and for all  $h \in [H]$ ,
$$
\mathcal {F} _ {h} ^ {\operatorname {l i n}} = \left\{f (z \mid x, a) = \left\langle \phi (x, a), w (z) \right\rangle : \quad \phi \in \Phi , w: [ 0, 1 ] \rightarrow \mathbb {R} ^ {d}, \right. \tag {2}
$$
$$
\left. \text {s . t .} \max  _ {z} \| w (z) \| _ {2} \leq \alpha \sqrt {d} \text {a n d} \max  _ {x, a, z} \left\langle \phi (x, a), w (z) \right\rangle \leq \alpha \right\},
$$
where  $\alpha := \max_{h,\pi,z,x,a} Z_h^\pi(z \mid x, a)$  is the largest mass for the cost-to-go distributions. In Appendix D, we show that  $\mathcal{F}^{\mathrm{lin}}$  satisfies DistBC. Further, if costs are discretized into a uniform grid of  $M$  points, its bracketing entropy is bounded by  $\widetilde{\mathcal{O}}(dM + \log |\Phi|)$ . Discretization is necessary to bound the statistical complexity of  $\mathcal{F}^{\mathrm{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  $\delta \in (0,1)$ , w.p. at least  $1 - \delta$ , running O-DISCO with UAE=TRUE and with  $\mathcal{F}^{\mathrm{lin}}$  as described above learns a policy  $\bar{\pi}$  such that,
$$
V ^ {\bar {\pi}} - V ^ {\star} \in \widetilde {\mathcal {O}} \left(H \sqrt {\frac {A d V ^ {\star} (d M + \log (| \Phi | / \delta))}{K}} + \frac {A d H ^ {2} (d M + \log (| \Phi | / \delta))}{K}\right).
$$
Proof. As described above, we have  $\mathrm{DE}_1(1 / K) \leq \mathcal{O}(d\log (dK))$  and  $\beta = \log (HK / \delta) + dM + \log |\Phi|$ . Since  $\mathrm{DistBC}$  is satisfied by  $\mathcal{F}^{\mathrm{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^{\star} \leq \widetilde{\mathcal{O}}(1/K)$ , O-DISCO has a fast  $\widetilde{\mathcal{O}}(1/K)$  convergence rate which improves over all prior results that converge at a slow  $\widetilde{\Omega}(1/\sqrt{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  $\mathcal{F}_k$ : (i)  $Z^{\star} \in \mathcal{F}_k$ , and (ii) elements of  $\mathcal{F}_k$  almost satisfy the distributional Bellman equation, i.e., for all  $h \in [H]$ , we have  $\sum_{i=1}^{k} \mathbb{E}_{\pi^i}[\delta_{h,k}(x_h, a_h)] \leq \mathcal{O}(\beta)$  where  $\delta_{h,k}(x_h, a_h) = D_\triangle(f_h^{(k)}(x_h, a_h)) \parallel T_h^{\star,D}f_{h+1}^{(k)}(x_h, a_h)$ . Next, we derive a corollary of Eq. ( $\triangle_1$ ):
$$
\left| \bar {f} - \bar {g} \right| \leq \sqrt {4 \bar {g} + D _ {\triangle} (f \| g)} \cdot \sqrt {D _ {\triangle} (f \| g)}. \tag {△2}
$$
To see why this is true, apply AM-GM to Eq.  $(\triangle_1)$  to get  $2(\bar{f} -\bar{g})\leq \bar{f} +\bar{g} +D_{\triangle}(f\parallel g)$ , which simplifies to  $\bar{f}\leq 3\bar{g} +D_{\triangle}(f\parallel g)$ . Plugging this back into Eq.  $(\triangle_1)$  yields Eq.  $(\triangle_2)$ . Then, by iterating Eq.  $(\triangle_2)$  and AM-GM, we derive a self-bounding lemma: for any  $f,\pi ,h$  we have  $\bar{f}_h(x_h,a_h)\lesssim Q_h^\pi (x_h,a_h) + H\sum_{t = h}^H\mathbb{E}_{\pi ,x_h,a_h}[D_\triangle (f_t(x_t,a_t)\parallel \mathcal{T}_h^{\pi ,D}f_{t + 1}(x_t,a_t))]$  (Lemma H.3).
Since  $\mathcal{T}_h^{\pi^k}\bar{f}_{h + 1}^{(k)}(x,a) = \overline{\mathcal{T}_h^{\pi^k,D}f_{h + 1}^{(k)}(x,a)}$  and  $\mathcal{T}_h^{\pi^k,D}f_{h + 1}^{(k)} = \mathcal{T}_h^{\star ,D}f_{h + 1}^{(k)}$ , we have
$$
\begin{array}{l} V ^ {\pi^ {k}} - V ^ {\star} \leq V ^ {\pi^ {k}} - \bar {f} _ {1} ^ {(k)} \left(x _ {1}, \pi_ {1} ^ {k} \left(x _ {1}\right)\right) \quad (\text {o p t i m i s m f r o m f a c t (i)}) \\ = \sum_ {h = 1} ^ {H} \mathbb {E} _ {\pi^ {k}} \left[ \mathcal {T} _ {h} ^ {\pi^ {k}} \bar {f} _ {h + 1} ^ {(k)} \left(x _ {h}, a _ {h}\right) - \bar {f} _ {h} ^ {(k)} \left(x _ {h}, a _ {h}\right) \right] \quad \text {(p e r f o r m a n c e d i f f e r e n c e)} \\ \leq 2 \sum_ {h = 1} ^ {H} \sqrt {\mathbb {E} _ {\pi^ {k}} [ \bar {f} _ {h} ^ {(k)} (x _ {h} , a _ {h}) + \delta_ {h , k} (x _ {h} , a _ {h}) ]} \sqrt {\mathbb {E} _ {\pi^ {k}} [ \delta_ {h , k} (x _ {h} , a _ {h}) ]} \quad (\text {E q .} (\triangle_ {2})) \\ \lesssim \sqrt {V ^ {\pi^ {k}} w + H \sum_ {h = 1} ^ {H} \mathbb {E} _ {\pi^ {k}} [ \delta_ {h , k} (x _ {h} , a _ {h}) ]} \sqrt {H \mathbb {E} _ {\pi^ {k}} [ \delta_ {h , k} (x _ {h} , a _ {h}) ]}. \quad (\text {L e m m a H . 3}) \\ \end{array}
$$
The implicit inequality  $V^{\pi^k} - V^{\star} \lesssim \sqrt{V^{\star} + H\sum_{h = 1}^{H}\mathbb{E}_{\pi^k}[\delta_{h,k}(x_h,a_h)]}\sqrt{H\mathbb{E}_{\pi^k}[\delta_{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  $\sum_{k = 1}^{K}\mathbb{E}_{\pi^k}[\delta_{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 Pessimistic Distributional Confidence set Optimization (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  $\pi$  based on a similar log-likelihood thresholding procedure as in O-DISCO, where the difference is we now use data sampled from  $\mathcal{T}_h^{\pi ,D}f_{h + 1}$  instead of  $\mathcal{T}_h^{\star ,D}f_{h + 1}$ . Next, output the policy with the most pessimistic mean amongst all the confidence sets.
Algorithm 3 Pessimistic Distributional Confidence set Optimization (P-DISCO)
1: Input: datasets  $\mathcal{D}_1, \ldots, \mathcal{D}_H$ , distribution function class  $\mathcal{F}$ , threshold  $\beta$ , policy class  $\Pi$ .
2: For all  $(h, f, \pi) \in [H] \times \mathcal{F} \times \Pi$ , sample  $y_{h,i}^{f,\pi} \sim f_{h+1}(x_{h,i}'', \pi_{h+1}(x_{h,i}')$ ), where  $(x_{h,i}, a_{h,i}, c_{h,i}, x_{h,i}')$  is the  $i$ -th datapoint of  $\mathcal{D}_h$ . Then, set  $z_{h,i}^{f,\pi} = c_{h,i} + y_{h,i}^{f,\pi}$  and define the confidence set,
$$
\mathcal {F} _ {\pi} = \left\{f \in \mathcal {F}: \sum_ {i = 1} ^ {N} \log f _ {h} \left(z _ {h, i} ^ {f, \pi} \mid x _ {h, i}, a _ {h, i}\right) \geq \max  _ {g \in \mathcal {F} _ {h}} \sum_ {i = 1} ^ {N} \log g \left(z _ {h, i} ^ {f, \pi} \mid x _ {h, i}, a _ {h, i}\right) - 7 \beta , \forall h \in [ H ] \right\}.
$$
3: For each  $\pi \in \Pi$ , define the pessimistic estimate  $f^{\pi} = \arg \max_{f\in \mathcal{F}_{\pi}}\mathbb{E}_{a\sim \pi (x_1)}\big[\bar{f}_1(x_1,a)\big]$ .
4: Output:  $\widehat{\pi} = \arg \max_{\pi \in \Pi} \mathbb{E}_{a \sim \pi(x_1)} \left[ \bar{f}_1^\pi(x_1, \pi) \right]$ .
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  $\widetilde{\pi}$ , with sub-optimality of the learned  $\widehat{\pi}$  degrading gracefully as coverage worsens. The coverage is measured by the single-policy concentrability  $C^{\widetilde{\pi}} = \max_{h} \left\| \frac{d d_{h}^{\widetilde{\pi}}}{d \nu_{h}} \right\|_{\infty}$ . 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  $\delta \in (0,1)$ , w.p. at least  $1 - \delta$ , running P-DISCO with  $\beta = \log(H|\Pi||\mathcal{F}| / \delta)$  learns a policy  $\widehat{\pi}$  that enjoys the following PAC bound with respect to any comparator policy  $\widetilde{\pi} \in \Pi$ :
$$
V ^ {\widetilde {\pi}} - V ^ {\widetilde {\pi}} \leq 9 H \sqrt {\frac {C ^ {\widetilde {\pi}} V ^ {\widetilde {\pi}} \beta}{N}} + \frac {3 0 H ^ {2} C ^ {\widetilde {\pi}} \beta}{N}.
$$
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^{\widetilde{\pi}}$  but also the comparator policy's value  $V^{\widetilde{\pi}}$ . In particular, P-DISCO can strongly compete with a comparator policy  $\widetilde{\pi}$  if one of the following is true: (i)  $\nu$  has good coverage over  $\widetilde{\pi}$ , so the  $\mathcal{O}(1 / \sqrt{N})$  term is manageable; or (ii)  $\widetilde{\pi}$  has small-loss, in which case we may even obtain a fast  $\mathcal{O}(1 / N)$  rate. Thus, P-DISCO has two chances at strongly competing with  $\widetilde{\pi}$ , 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-of-the-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.
Algorithm:	SquareCB	FastCB	DistCB (Ours)
King County Housing [Vanschoren et al., 2013]
All episodes	.756 (.0007)	.734 (.0007)	.726 (.0003)
Last 100 ep.	.725 (.0012)	.719 (.0013)	.708 (.0019)
Prudential Life Insurance [Montoya et al., 2015]
All episodes	.456 (.0082)	.491 (.0029)	.411 (.0038)
Last 100 ep.	.481 (.0185)	.474 (.0111)	.388 (.0086)
CIFAR-100 [Krizhevsky, 2009]
All episodes	.872 (.0010)	.856 (.0016)	.838 (.0021)
Last 100 ep.	.828 (.0024)	.793 (.0031)	.775 (.0027)
Table 1: Avg cost over all episodes and last 100 episodes (lower is better). We report 'mean (sem)' over 10 seeds.
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 \times (y - \hat{y})$  where  $y$  is the actual risk level, and  $\hat{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.