SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Last step: bounding  $\sum_{k} \Delta_{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  $\varepsilon = 1 / K$ .
On one hand, if UAE  $\equiv$  FALSE, then,
$$
\sum_ {k = 1} ^ {K} \Delta_ {k} = \sum_ {h = 1} ^ {H} \sum_ {k = 1} ^ {K} \mathbb {E} _ {\pi^ {k}} \left[ \delta_ {h, k} \left(x _ {h}, a _ {h}\right) \right] \leq 1 0 0 0 H \operatorname {D E} _ {1} (1 / K) \beta \log (K).
$$
On the other hand, if UAE  $\equiv$  TRUE, then, we use the V-type analogs,
$$
\begin{array}{l} \sum_ {k = 1} ^ {K} \Delta_ {k} = \sum_ {h = 1} ^ {H} \sum_ {k = 1} ^ {K} \mathbb {E} _ {\pi^ {k}} [ \delta_ {h, k} (x _ {h}, a _ {h}) ] \\ \leq A \sum_ {h = 1} ^ {H} \sum_ {k = 1} ^ {K} \mathbb {E} _ {\pi^ {k}} \left[ \mathbb {E} _ {a \sim \operatorname {u n i f} (\mathcal {A})} \delta_ {h, k} \left(x _ {h}, a\right) \right] \\ \leq 1 0 0 0 A H \mathrm {D E} _ {1} (1 / K) \beta \log (K). \\ \end{array}
$$
This concludes the proof for both the regret and PAC bounds.
Lemma H.3 (Self-bounding lemma). Let  $f \in \mathcal{F}$  and let  $\pi$  be any policy. Let us denote  $\delta_h(x, a) \coloneqq D_{\triangle}(f_h(x, a) \parallel \mathcal{T}_h^{\pi, D}f_{h+1}(x, a))$ . Then, for all  $h \in [H]$ , for all  $x_h, a_h$ , we have
$$
\bar {f} _ {h} \left(x _ {h}, a _ {h}\right) \leq e Q _ {h} ^ {\pi} \left(x _ {h}, a _ {h}\right) + 4 H \sum_ {t = h} ^ {H} \mathbb {E} _ {\pi , x _ {h}, a _ {h}} \left[ \delta_ {t} \left(x _ {t}, a _ {t}\right) \right].
$$
Proof. We prove the following refined subclaim inductively: for all  $h \in [H]$ , for all  $x_h, a_h$ , we have
$$
\bar {f} _ {h} \left(x _ {h}, a _ {h}\right) \leq \sum_ {t = h} ^ {H} \left(1 + \frac {1}{\bar {H}}\right) ^ {t - h} \mathbb {E} _ {\pi , x _ {h}, a _ {h}} \left[ \bar {c} _ {t} \left(x _ {t}, a _ {t}\right) + 2 H \delta_ {t} \left(x _ {t}, a _ {t}\right) \right]. \tag {IH}
$$
For  $H + 1$  this is trivially true. Now fix any  $h$  and suppose IH is true for  $h + 1$ . By Eq.  $(\triangle_2)$ , for any  $h, x_h, a_h$ , we have,
$$
\begin{array}{l} \bar {f} _ {h} \left(x _ {h}, a _ {h}\right) - \mathcal {T} _ {h} ^ {\pi} \bar {f} _ {h + 1} \left(x _ {h}, a _ {h}\right) \leq \sqrt {4 \mathcal {T} _ {h} ^ {\pi} \bar {f} _ {h + 1} \left(x _ {h} , a _ {h}\right) + \delta_ {h} \left(x _ {h} , a _ {h}\right)} \sqrt {\delta_ {h} \left(x _ {h} , a _ {h}\right)} \\ \leq \sqrt {4 \mathcal {T} _ {h} ^ {\pi} \bar {f} _ {h + 1} (x _ {h} , a _ {h}) \delta_ {h} (x _ {h} , a _ {h})} + \delta_ {h} (x _ {h}, a _ {h}) \\ \leq \frac {1}{H} \mathcal {T} _ {h} ^ {\pi} \bar {f} _ {h + 1} \left(x _ {h}, a _ {h}\right) + (H + 1) \delta_ {h} \left(x _ {h}, a _ {h}\right). \quad (\mathrm {A M - G M}) \\ \end{array}
$$
In particular, we have that
$$
\begin{array}{l} \bar {f} _ {h} \left(x _ {h}, a _ {h}\right) \\ \leq \left(1 + \frac {1}{H}\right) \mathcal {T} _ {h} ^ {\pi} \bar {f} _ {h + 1} \left(x _ {h}, a _ {h}\right) + 2 H \delta_ {h} \left(x _ {h}, a _ {h}\right) \\ = \left(1 + \frac {1}{H}\right) \left(\bar {c} _ {h} \left(x _ {h}, a _ {h}\right) + \mathbb {E} _ {x _ {h + 1} \sim P _ {h} ^ {\star} \left(x _ {h}, a _ {h}\right)} \left[ \bar {f} _ {h + 1} \left(x _ {h + 1}, \pi\right) \right]\right) + 2 H \delta_ {h} \left(x _ {h}, a _ {h}\right) \\ \leq \left(1 + \frac {1}{H}\right) \binom {\bar {c} _ {h} \left(x _ {h}, a _ {h}\right) + \mathbb {E} _ {x _ {h + 1} \sim P _ {h} ^ {*} \left(x _ {h}, a _ {h}\right)} \left[ \sum_ {t = h + 1} ^ {H} \left(1 + \frac {1}{H}\right) ^ {t - h - 1} \mathbb {E} _ {\pi , x _ {h + 1}} \left[ \bar {c} _ {t} \left(x _ {t}, a _ {t}\right) + 2 H \delta_ {t} \left(x _ {t}, a _ {t}\right)\right]\right]\left. \right) \\ + 2 H \delta_ {h} \left(x _ {h}, a _ {h}\right), \\ \end{array}
$$
which proves the inductive claim. Noting that  $\sum_{t=1}^{H} (1 + 1 / H)^{t} \leq 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  $\delta \in (0,1)$  and set  $\beta = \log(HK|\mathcal{F}| / \delta)$ . Then, w.p. at least  $1 - \delta$ ,
$$
\operatorname {R e g r e t} _ {\mathrm {O - D I S C O}} (K) \in \mathcal {O} \left(H \sqrt {X A K V ^ {\star} \beta} + H ^ {2} X A \beta\right).
$$
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  $\widetilde{\mathcal{O}}(\sqrt{XAK})$ . The key benefit over prior bounds is that our leading term scales with the minimum cost of the problem  $V^{\star}$ . For example, if  $V^{\star} \approx 0$ , O-DISCO attains  $\mathcal{O}(\log K)$  regret while uniform regret bounds are lower bounded by  $\Omega(\sqrt{K})$ . Compared to the minimax-optimal UCBVI [Azar et al., 2017], one weakness of our theorem is that it needs a  $\mathcal{F}$  satisfying BC. Fortunately, in tabular MDPs where cost is only revealed at the last step from a known distribution, we can choose  $\mathcal{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  $\mathcal{F}_{tab}$  yields  $\beta = \mathcal{O}(X^{2}A^{2}\log (XAHK / \delta))$ . We note that if costs are unknown but discrete, it is possible to construct a BC function class with  $\beta$  scaling as  $\mathcal{O}(X^{2}A^{2}\log (nXAHK / \delta))$  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  $\phi^{\star}(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}^{\pi}(z\mid x,a)$  is lower bounded, i.e., the value distribution is sufficiently smooth.
I Proofs for Offline RL
$$
V ^ {\widehat {\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}.
$$
Proof of Theorem 6.1. For shorthand, let  $\delta_h^\pi (x,a) = D_\triangle (f_h^\pi (x,a)\parallel \mathcal{T}_h^{\pi ,D}f_{h + 1}^\pi (x,a))$  and  $\Delta^{\pi} = \sum_{h = 1}^{H}\mathbb{E}_{\pi}[\delta_h^\pi (x_h,a_h)]$ . Also, let  $f(x,\pi) = \mathbb{E}_{a\sim \pi (x)}[f(x,a)]$ .
By Theorem F.3, we have the following two facts, for all  $\pi \in \Pi$ ,
(i) Pessimism:  $V^{\pi} \leq \bar{f}_1^{\pi}(x_1, \pi)$  (since  $Z^{\pi} \in \mathcal{F}_{\pi}$ ) for all  $\pi \in \Pi$ , and
(ii)  $\mathbb{E}_{\nu_h}[\delta_h^\pi (x_h,a_h)]\stackrel {\sim}{\leq}\beta 'N^{-1}$  for all  $h$  where Theorem F.3 and the fact that  $D_{\triangle}\leq 4H^{2}$  certifies that  $\beta^{\prime} = 240\beta$  is sufficient.
With these two facts, we can bound the suboptimality of  $\widehat{\pi}$  as follows:
$$
\begin{array}{l} V ^ {\widehat {\pi}} - V ^ {\widetilde {\pi}} \\ \leq \bar {f} _ {1} ^ {\hat {\pi}} \left(x _ {1}, \hat {\pi}\right) - V ^ {\hat {\pi}} (Fact(i)) \\ \leq \bar {f} _ {1} ^ {\pi} \left(x _ {1}, \widetilde {\pi}\right) - V ^ {\widetilde {\pi}} \quad \text {(P o l i c y s e l e c t i o n s c h e m e i n A l g o r a t h m 3 (L i n e 4))} \\ = \sum_ {h = 1} ^ {H} \mathbb {E} _ {\widetilde {\pi}} \left[ \widetilde {f} _ {h} ^ {\widetilde {\pi}} \left(x _ {h}, \widetilde {\pi}\right) - \mathcal {T} _ {h} ^ {\widetilde {\pi}} \widetilde {f} _ {h + 1} ^ {\widetilde {\pi}} \left(x _ {h}, a _ {h}\right) \right] (PDLLemmaH.2) \\ \leq \sum_ {h = 1} ^ {H} \sqrt {\mathbb {E} _ {\widetilde {\pi}} \left[ 4 f _ {h} ^ {\widetilde {\pi}} \left(x _ {h} , a _ {h}\right) + \delta_ {h} ^ {\widetilde {\pi}} \left(x _ {h} , a _ {h}\right) \right]} \sqrt {\mathbb {E} _ {\widetilde {\pi}} \left[ \delta_ {h} ^ {\widetilde {\pi}} \left(x _ {h} , a _ {h}\right) \right]} (Eq. (△2)) \\ \leq \sum_ {h = 1} ^ {H} \sqrt {4 e V ^ {\tilde {\pi}} + 1 7 H \sum_ {t = h} ^ {H} \mathbb {E} _ {\tilde {\pi}} \left[ \delta_ {t} ^ {\tilde {\pi}} \left(x _ {t} , a _ {t}\right) \right]} \sqrt {\mathbb {E} _ {\tilde {\pi}} \left[ \delta_ {h} ^ {\tilde {\pi}} \left(x _ {h} , a _ {h}\right) \right]} (Lemma H.3) \\ \leq \sqrt {4 e V ^ {\widetilde {\pi}} + 1 7 H \Delta^ {\widetilde {\pi}}} \sqrt {H \Delta^ {\widetilde {\pi}}} \\ \leq 4 \sqrt {H V ^ {\widetilde {\pi}} \Delta^ {\widetilde {\pi}}} + 5 H \Delta^ {\widetilde {\pi}}. \\ \end{array}
$$
Finally, we can bound  $\Delta^{\widetilde{\pi}}$  by a change of measure,
$$
\begin{array}{l} \Delta^ {\widetilde {\pi}} = \sum_ {h = 1} ^ {H} \mathbb {E} _ {\widetilde {\pi}} \left[ \delta_ {h} ^ {\widetilde {\pi}} \left(x _ {h}, a _ {h}\right) \right] \\ \leq C ^ {\widetilde {\pi}} \sum_ {h = 1} ^ {H} \mathbb {E} _ {\nu_ {h}} [ \delta_ {h} (x _ {h}, a _ {h}) ] \\ \leq C ^ {\pi} H \cdot \beta^ {\prime} N ^ {- 1}. \tag {Fact(ii)} \\ \end{array}
$$
$$
V ^ {\widehat {\pi}} - V ^ {\widetilde {\pi}} \leq 4 H \sqrt {\frac {C ^ {\widetilde {\pi}} V ^ {\widetilde {\pi}} \beta^ {\prime}}{N}} + \frac {5 H ^ {2} C ^ {\widetilde {\pi}} \beta^ {\prime}}{N}.
$$
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  $\delta \in (0,1)$  and set  $\beta = \log(HK|\mathcal{F}| / \delta)$  and  $\beta' = 60\beta$ . Then, w.p. at least  $1 - \delta$ , running O-DISCO (Algorithm 4) with UAE = FALSE yields the following small-loss regret bound,
$$
\operatorname {R e g r e t} _ {\mathrm {O - D I S C O}} (K) \leq 5 H \sqrt {K V ^ {\star} \operatorname {L S E C} (K) \beta^ {\prime}} + 1 8 H ^ {2} \operatorname {L S E C} (K) \beta^ {\prime}. \tag {12}
$$
If instead UAE = TRUE, the outputted policy  $\bar{\pi}$  enjoys the following small-loss PAC bound,
$$
V ^ {\star} - V ^ {\bar {\pi}} \leq 5 H \sqrt {\frac {A V ^ {\star} \mathrm {L S E C} _ {v} (K) \beta^ {\prime}}{K}} + 1 8 H ^ {2} \frac {A \mathrm {L S E C} _ {v} (K) \beta^ {\prime}}{K}.
$$
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 \in [K]$ ,
(i) Optimism:  $V^{\star} \leq \max_{a} \bar{f}_{1}^{(k)}(x_{1}, a)$  (since  $Z^{\star} \in \mathcal{F}_{k}$ ) and
(ii)  $\sum_{i < k} \mathbb{E}_{\pi^i}[\delta_{h,k}(s_h, a_h)] \leq \beta'$  for all  $h$ . If  $\mathrm{UAE} = \mathrm{TRUE}$ , then  $a_h$  is sampled from  $\mathrm{unif}(\mathcal{A})$  rather than  $\pi^i$ , i.e., we have  $\sum_{i < k} \mathbb{E}_{s_h \sim \pi^i, a_h \sim \mathrm{unif}(\mathcal{A})}[\delta_{h,k}(s_h, a_h)] \leq \beta'$ , where  $\beta' \lesssim \beta$ . Theorem F.2 certifies that  $\beta' = 60\beta$  is sufficient.
Fix any episode  $k \in [K]$ . Then,
$$
\begin{array}{l} V ^ {\star} - V ^ {\pi^ {k}} \\ \leq \max  _ {a} \bar {f} _ {1} ^ {(k)} \left(x _ {1}, a\right) - V ^ {\pi^ {k}} (Fact(i)) \\ = \sum_ {h = 1} ^ {H} \mathbb {E} _ {\pi^ {k}} \left[ \bar {f} _ {h} ^ {(k)} \left(x _ {h}, \pi_ {h} ^ {k} \left(x _ {h}\right)\right) - \mathcal {T} _ {h} ^ {\pi^ {k}} \bar {f} _ {h + 1} ^ {(k)} \left(x _ {h}, a _ {h}\right) \right] (PDLLemmaH.2) \\ = \sum_ {h = 1} ^ {H} \mathbb {E} _ {\pi^ {k}} \left[ \bar {f} _ {h} ^ {(k)} \left(x _ {h}, a _ {h}\right) - \overline {{\mathcal {T} _ {h} ^ {\pi^ {k} , D} f _ {h + 1} ^ {(k)}}} \left(x _ {h}, a _ {h}\right) \right] (Lemma H.1) \\ \leq \sum_ {h = 1} ^ {H} \sqrt {\mathbb {E} _ {\pi^ {k}} \left[ 4 \bar {f} _ {h} ^ {(k)} \left(x _ {h} , a _ {h}\right) + \delta_ {h , k} \left(x _ {h} , a _ {h}\right) \right]} \cdot \sqrt {\mathbb {E} _ {\pi^ {k}} \left[ \delta_ {h , k} \left(x _ {h} , a _ {h}\right) \right]} \quad (\text {E q .} (\triangle_ {2})) \\ \leq \sum_ {h = 1} ^ {H} \sqrt {4 e V ^ {\pi^ {k}} + 1 7 H \sum_ {t = h} ^ {H} \mathbb {E} _ {\pi^ {k}} [ \delta_ {t , k} (x _ {t} , a _ {t}) ]} \cdot \sqrt {\mathbb {E} _ {\pi^ {k}} [ \delta_ {h , k} (x _ {h} , a _ {h}) ]} \\ \leq \sqrt {4 e V ^ {\pi^ {k}} + 1 7 H \Delta_ {k}} \cdot \sqrt {H \Delta_ {k}} (♣) \\ \leq \sqrt {4 e V ^ {\star} + 1 7 H \Delta_ {k}} \cdot \sqrt {H \Delta_ {k}} \\ \end{array}
$$
$$
\begin{array}{l} \sum_ {k = 1} ^ {K} V ^ {\pi^ {k}} - V ^ {\star} \leq \sum_ {k = 1} ^ {K} \sqrt {4 e V ^ {\star} + 1 7 H \Delta_ {k}} \sqrt {H \Delta_ {k}} \\ \leq \sqrt {4 e K V ^ {\star} + 1 7 H \sum_ {k} \Delta_ {k}} \sqrt {H \sum_ {k} \Delta_ {k}} \quad (\text {C a u c h y - S c h w a r t z}) \\ \leq 5 \sqrt {H K V ^ {\star} \sum_ {k} \Delta_ {k}} + 1 8 H \sum_ {k} \Delta_ {k}. \\ \end{array}
$$
The bounds for  $\Delta_{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  $\clubsuit$ . However, in the reward-maximizing setting, we simply bound  $V^{\pi^k} \leq V^\star$ , 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  $\delta \in (0,1)$  and set  $\beta = \log (H|\Pi ||\mathcal{F}| / \delta)$ . Then, w.p. at least  $1 - \delta$ , P-DISCO (Algorithm 4) learns a policy  $\widehat{\pi}$  such that for any comparator policy  $\widetilde{\pi} \in \Pi$ , we have
$$
V ^ {\widetilde {\pi}} - V ^ {\widehat {\pi}} \leq 9 H \sqrt {\frac {C ^ {\widetilde {\pi}} V ^ {\widetilde {\pi}} \beta}{N}} + \frac {3 0 H ^ {2} C ^ {\widetilde {\pi}} \beta}{N}.
$$
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  $\pi \in \Pi$ ,
(i) Pessimism:  $\tilde{f}_1^\pi(x_1, \pi) \leq V^\pi$  (since  $Z^\pi \in \mathcal{F}_\pi$ ) for all  $\pi \in \Pi$ , and  
(ii)  $\mathbb{E}_{\nu_h}[\delta_h^\pi(x_h, a_h)] \leq \beta' N^{-1}$  for all  $h$  where  $\beta' \leq 60\beta$ .
$$
\begin{array}{l} V ^ {\widetilde {\pi}} - V ^ {\widehat {\pi}} \\ \leq V ^ {\bar {\pi}} - \bar {f} _ {1} ^ {\bar {\pi}} \left(x _ {1}, \widehat {\pi}\right) (Fact(i)) \\ \leq V ^ {\widetilde {\pi}} - \tilde {f} _ {1} ^ {\widetilde {\pi}} \left(x _ {1}, \widetilde {\pi}\right) \quad \text {(P o l i c y s e l e c t i o n r u l e i n L i n e 5)} \\ = \sum_ {h = 1} ^ {H} \mathbb {E} _ {\widetilde {\pi}} \left[ \mathcal {T} _ {h} ^ {\widetilde {\pi}} \bar {f} _ {h + 1} ^ {\widetilde {\pi}} \left(x _ {h}, a _ {h}\right) - \bar {f} _ {h} ^ {\widetilde {\pi}} \left(x _ {h}, \widetilde {\pi}\right) \right] (PDLLemmaH.2) \\ \leq \sum_ {h = 1} ^ {H} \sqrt {\mathbb {E} _ {\widetilde {\pi}} \left[ 4 \bar {f} _ {h} ^ {\widetilde {\pi}} \left(x _ {h} , a _ {h}\right) + \delta_ {h} ^ {\widetilde {\pi}} \left(x _ {h} , a _ {h}\right) \right]} \sqrt {\mathbb {E} _ {\widetilde {\pi}} \left[ \delta_ {h} ^ {\widetilde {\pi}} \left(x _ {h} , a _ {h}\right) \right]}. (Eq. (△2)) \\ \end{array}
$$
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  $\gamma$  learning rate at each time step  $t$  is set to  $\gamma_{t} = \gamma_{0}t^{p}$  where  $\gamma_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.
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.
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.
Datasets
Dataset	Items	Number of Actions	Number of Costs
CIFAR-100	50,000	100	3
Prudential Life Insurance	59,381	8	9
King County Housing	20,148	100	101
Table 3: Overview of the three datasets and their experimental setups
Prudential Life Insurance This dataset is from the Prudential Life Insurance Kaggle competition [Montoya et al., 2015]. It 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 - \hat{y})$  where  $y$  is the actual risk level, and  $\hat{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 DISTRCB 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.
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.