SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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^{\star}$ , 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.
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 distributional-fitting scheme to inspire a practical offline RL algorithm with small loss guarantees, à la Cheng et al. [2022]. Finally, it would be interesting to investigate other loss functions that yield small-loss 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.
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:13626-13640, 2021.
Tianjun Zhang, Tongzheng Ren, Mengjiao Yang, Joseph Gonzalez, Dale Schuurmans, and Bo Dai. Making linear mdps practical via contrastive representation learning. In International Conference on Machine Learning, pages 26447-26466. PMLR, 2022.
Tong Zhang. Mathematical Analysis of Machine Learning Algorithms. 2023. http://www.tongzhang-ml.org/1t-book.html.
Zhengyuan Zhou, Susan Athey, and Stefan Wager. Offline multi-action policy learning: Generalization and optimization. Operations Research, 71(1):148-183, 2023.
Table 2: List of Notations
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 = Ey~d[y].
[N]	{1,2,...,N} for any natural number N.
Zh(x,a)	Distribution of ∑Ht=hct given xh=x,ah=a rolling in from π.
Qh(x,a),Vh(x)	Qh(x,a) = Zh(x,a) and Vh= Ea~π(x)[Qh(x,a)].
π*	Optimal policy, i.e., π* = arg minπV1(𝑥1).
Without loss of optimality, we take π*: X → A to be Markov & deterministic.
Zh*,Qh*,Vh*	Zh*,Qh*,Vh* with π = π*, the optimal policy.
Th*,Th*	The Bellman operators that act on functions.
Th*,D,Th*,D	The distributional Bellman operators that act on conditional distributions.
Vπ,Zπ,V*,Z*	Vπ = V1(𝑥1), Zπ = Z1(𝑥1). V*, Z* are defined similarly with π*.
dh(x,a)	The probability of π visiting (x,a) at time h.
Cπ	Coverage coefficient maxh||ddh/duh||∞.
D△(f || g)	Triangular discrimination between f, g.
H(f || g)	Hellinger distance between f, g.
DKL(f || g)	KL divergence between f, g.
A.1 Statistical Distances
Let  $f, g$  be distributions over  $\mathcal{V}$ . Then,
$$
D _ {\triangle} (f \parallel g) = \sum_ {y} \frac {(f (y) - g (y)) ^ {2}}{f (y) + g (y)},
$$
$$
H (f \parallel g) = \sqrt {\frac {1}{2} \sum_ {y} \left(\sqrt {f (y)} - \sqrt {g (y)}\right) ^ {2}},
$$
$$
D _ {K L} (f \parallel g) = \sum_ {y} f (y) \log (f (y) / g (y)),
$$
$$
D _ {T V} (f \parallel g) = \frac {1}{2} \sum_ {y} | f (y) - g (y) |.
$$
The following standard inequalities will be helpful:
$$
H ^ {2} \leq D _ {T V} \leq \sqrt {2} H,
$$
$$
2 H ^ {2} \leq D _ {\triangle} \leq 4 H ^ {2}, \tag {Lemma A.1}
$$
$$
H \leq \sqrt {D _ {K L}}.
$$
Lemma A.1. For any distributions  $f, g$ , we have  $2H^{2}(f \parallel g) \leq D_{\triangle}(f \parallel g) \leq 4H^{2}(f \parallel g)$ .
$$
D _ {\triangle} (f \parallel g) = \int_ {y} \left(\frac {f (y) - g (y)}{\sqrt {f (y) + g (y)}}\right) ^ {2}.
$$
Applying  $\frac{1}{\sqrt{f(y) + \sqrt{g(y)}}} \leq \frac{1}{\sqrt{f(y) + g(y)}} \leq \frac{\sqrt{2}}{\sqrt{f(y) + \sqrt{g(y)}}}$  concludes the proof.
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  $\mathcal{F}$ , threshold  $\beta$ , flag UAE, flag SMALLRETURN.  
2: Initialize  $\mathcal{D}_{h,0} \gets \emptyset$  for all  $h \in [H]$ , and set  $\mathcal{F}_0 = \mathcal{F}$ .  
3: Set op = max if SMALLRETURN else op = min.  
4: for episode  $k = 1, 2, \ldots, K$  do  
5: Set  $f^{(k)} = \arg \mathrm{op}_{f \in \mathcal{F}_{k-1}} \mathrm{op}_a \bar{f}_1(x_1, a)$ .  
6: Set  $\pi_h^k(x) = \arg \mathrm{op}_a \bar{f}_h^{(k)}(x, a)$ .  
7: if UAE then  
8: For each  $h \in [H]$ , collect  $x_{h,k} \sim d_h^{\pi^k}$ ,  $a_{h,k} \sim \mathrm{unif}(\mathcal{A})$ ,  $c_{h,k} \sim C_h(x_{h,k}, a_{h,k})$ ,  $x_{h,k}' \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}')\}$ .  
9: else  
10: 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})\}$ .  
11: end if  
12: For all  $(h,f) \in [H] \times \mathcal{F}$ , sample  $y_{h,i}^f \sim f_{h+1}(x_{h,i}', a')$  and  $a' = \arg \mathrm{op}_a \bar{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  $\mathcal{D}_{h,k}$ . Also, 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(z_{h,i}^f | x_{h,i}, a_{h,i}) \geq \max_{\tilde{f} \in \mathcal{F}} \sum_{i=1}^k \log \widetilde{f}_h(z_{h,i}^f | x_{h,i}, a_{h,i}) - 7\beta, \forall h \in [H] \right\}. ]  
13: end for  
14: Output:  $\bar{\pi} = \mathrm{unif}(\pi^{1:K})$ .
Algorithm 5 P-DISCO (with small return)
1: Input: datasets  $\mathcal{D}_1, \ldots, \mathcal{D}_H$ , distribution function class  $\mathcal{F}$ , threshold  $\beta$ , policy class  $\Pi$ , flag SMALLRETURN.  
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(z_{h,i}^{f,\pi} \mid x_{h,i}, a_{h,i}) \geq \max_{\widetilde{f} \in \mathcal{F}} \sum_{i=1}^{N} \log \widetilde{f}_h(z_{h,i}^{f,\pi} \mid x_{h,i}, a_{h,i}) - 7\beta, \forall h \in [H] \right\}. ]  
3: Set op = max if SMALLRETURN else op = min.  
4: For each  $\pi \in \Pi$ , define the pessimistic estimate  $f^{\pi} = \arg \mathrm{op}_{f \in \mathcal{F}_{\pi}} \mathbb{E}_{a \sim \pi(x_1)}[\bar{f}_1(x_1, a)]$ .  
5: Output:  $\widehat{\pi} = \arg \mathrm{op}_{\pi \in \Pi} \mathbb{E}_{a \sim \pi(x_1)}[\bar{f}_1^{\pi}(x_1, \pi)]$ .
C Proofs for DistCB
Lemma C.1 (Azuma). Let  $\{X_i\}_{i \in [N]}$  be a sequence of random variables supported on  $[0, 1]$ , adapted to filtration  $\{\mathcal{F}_i\}_{i \in [N]}$ . For any  $\delta \in (0, 1)$ , we have w.p. at least  $1 - \delta$ ,
$$
\sum_ {t = 1} ^ {N} \mathbb {E} \left[ X _ {t} \mid \mathcal {F} _ {t - 1} \right] \leq \sum_ {t = 1} ^ {N} X _ {t} + \sqrt {N \log (2 / \delta)}, \quad \text {(S t a n d a r d A z u m a)}
$$
$$
\sum_ {t = 1} ^ {N} \mathbb {E} \left[ X _ {t} \mid \mathcal {F} _ {t - 1} \right] \leq 2 \sum_ {t = 1} ^ {N} X _ {t} + 2 \log (1 / \delta). \quad (\text {M u l t i p l i c a t i v e A z u m a})
$$
Proof. For standard Azuma, see Zhang [2023, Theorem 13.4]. For multiplicative Azuma, apply [Zhang, 2023, Theorem 13.5] with  $\lambda = 1$ . The claim follows, since  $\frac{1}{1 - \exp(-\lambda)} \leq 2$ .
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)$ ,
$$
\mathrm {R e g r e t} _ {\mathrm {D I S T C B}} (K) \leq 2 3 2 \sqrt {A C ^ {\star} \mathrm {R e g r e t} _ {\log} (K) \log (1 / \delta)} + 2 3 0 0 A \big (\mathrm {R e g r e t} _ {\log} (K) + \log (1 / \delta) \big).
$$
Proof of Theorem 4.1. First, recall the per-step inequality of ReIGW Foster and Krishnamurthy [2021, Theorem 4], which states: for any  $\widehat{f}$  and  $\gamma \geq 2A$ , if we set  $p = \mathrm{ReIGW}_{\gamma}(\widehat{f},\gamma)$ , then, for all  $f\in [0,1]^A$ , we have
$$
\sum_ {a} p (a) (f (a) - f (a ^ {\star})) \leq \frac {5 A}{\gamma} \sum_ {a} p (a) f (a) + 7 \gamma \sum_ {a} p (a) \frac {(\widehat {f} (a) - f (a)) ^ {2}}{\widehat {f} (a) + f (a)},
$$
where  $a^{\star} = \arg \min_{a} f(a)$ . For any  $k \in [K]$ , applying this to  $\widehat{f} = \bar{f}^{(k)}(s_k, \cdot)$ ,  $p = p_k$  and  $f = \bar{C}(s_k, \cdot)$ , we have
$$
\begin{array}{l} \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \bar {C} \left(s _ {k}, a _ {k}\right) - \bar {C} \left(s _ {k}, \pi^ {\star} \left(s _ {k}\right)\right) \right] \leq \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \frac {5 A}{\gamma} \bar {C} \left(s _ {k}, a _ {k}\right) + 7 \gamma \frac {\left(\bar {f} ^ {(k)} \left(s _ {k} , a _ {k}\right) - \bar {C} \left(s _ {k} , a _ {k}\right)\right) ^ {2}}{\bar {f} ^ {(k)} \left(s _ {k} , a _ {k}\right) + \bar {C} \left(s _ {k} , a _ {k}\right)} \right] \\ \leq \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \frac {5 A}{\gamma} \bar {C} \left(s _ {k}, a _ {k}\right) + 7 \gamma D _ {\triangle} \left(f ^ {(k)} \left(s _ {k}, a _ {k}\right) \| C \left(s _ {k}, a _ {k}\right)\right) \right] \tag {Eq. (△₁)} \\ \end{array}
$$
Since  $D_{\triangle} \leq 4H^{2}$ , we have
$$
\begin{array}{l} \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ D _ {\triangle} \left(f ^ {(k)} \left(s _ {k}, a _ {k}\right) \parallel C \left(s _ {k}, a _ {k}\right)\right) \right] \\ \leq 4 \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ H ^ {2} \left(C \left(s _ {k}, a _ {k}\right) \| f ^ {(k)} \left(s _ {k}, a _ {k}\right)\right) \right] \\ \leq 8 \sum_ {k = 1} ^ {K} H ^ {2} \left(C \left(s _ {k}, a _ {k}\right) \| f ^ {(k)} \left(s _ {k}, a _ {k}\right)\right) + 8 \log (1 / \delta) (\text {M u l t i p l i c a t i v e A z u m a , s i n c e} H ^ {2} \in [ 0, 1 ]) \\ \leq 8 \operatorname {R e g r e t} _ {\log} (K) + 1 0 \log (1 / \delta). \quad \text {(F o s t e r a l . [ 2 0 2 1 , L e m m a A . 1 4 ])} \\ \end{array}
$$
$$
\sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \bar {C} (s _ {k}, a _ {k}) - \bar {C} (s _ {k}, \pi^ {\star} (s _ {k})) \right] \leq \frac {5 A}{\gamma} \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \bar {C} (s _ {k}, a _ {k}) \right] + 7 0 \gamma \big (\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta) \big).
$$
Finally, recalling that  $1 / (1 - \varepsilon) \leq 1 + 2\varepsilon$  when  $\varepsilon \leq \frac{1}{2}$ , and the fact that  $\frac{5A}{\gamma} \leq \frac{1}{2}$ , we have
$$
\sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \bar {C} (s _ {k}, a _ {k}) - \bar {C} (s _ {k}, \pi^ {\star} (s _ {k})) \right] \leq \frac {1 0 A}{\gamma} \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \bar {C} (s _ {k}, \pi^ {\star} (s _ {k})) \right] + 1 4 0 \gamma \big (\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta) \big).
$$
By Azuma's inequality, we have
$$
\begin{array}{l} \sum_ {k = 1} ^ {K} \bar {C} (s _ {k}, a _ {k}) - \bar {C} (s _ {k}, \pi^ {\star} (s _ {k})) \\ \leq 2 \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \bar {C} \left(s _ {k}, a _ {k}\right) - \bar {C} \left(s _ {k}, \pi^ {\star} \left(s _ {k}\right)\right) \right] + 2 \log (1 / \delta) \\ \leq \frac {2 0 A}{\gamma} \sum_ {k = 1} ^ {K} \mathbb {E} _ {a _ {k}} \left[ \bar {C} \left(s _ {k}, \pi^ {\star} \left(s _ {k}\right)\right) \right] + 1 4 0 \gamma \left(\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta)\right) + 2 \log (1 / \delta) \\ \leq \frac {4 0 A}{\gamma} \left(C ^ {\star} + \log (1 / \delta)\right) + 1 4 0 \gamma \left(\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta)\right) + 2 \log (1 / \delta). \\ \end{array}
$$
(Multiplicative Azuma)
Now set  $\gamma = \sqrt{\frac{40A(C^{\star} + \log(1 / \delta))}{140\left(\mathrm{Regret}_{\log}(K) + \log(1 / \delta)\right)}}\lor 10A.$
Case 1 is when  $\sqrt{\frac{40A(C^{\star} + \log(1 / \delta))}{140\left(\mathrm{Regret}_{\log}(K) + \log(1 / \delta)\right)}} \leq 10A$ , i.e.,  $(C^{\star} + \log (1 / \delta)) \leq$
$280A\big(\mathrm{Regret}_{\log}(K) + \log(1/\delta)\big)$ , we have the above is at most
$$
\begin{array}{l} 4 \left(C ^ {\star} + \log (1 / \delta)\right) + 1 1 2 0 A \left(\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta)\right) + 2 \log (1 / \delta) \\ \leq 2 2 4 0 A \left(\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta)\right) + 2 \log (1 / \delta). \\ \end{array}
$$
Case 2 is when the left term dominates, then the bound is,
$$
\begin{array}{l} 2 \sqrt {4 4 8 0 A (C ^ {\star} + \log (1 / \delta)) \left(\operatorname {R e g r e t} _ {\log} (K) + \log (1 / \delta)\right)} + 2 \log (1 / \delta) \\ \leq 2 \sqrt {1 3 4 4 0 A C ^ {\star} \operatorname {R e g r e t} _ {\log} (K) \log (1 / \delta) + 4 4 8 0 A \log^ {2} (1 / \delta)} + 2 \log (1 / \delta) \\ \leq 2 3 2 \sqrt {A C ^ {\star} \operatorname {R e g r e t} _ {\log} (K) \log (1 / \delta)} + 1 3 4 \sqrt {A} \log (1 / \delta) + 2 \log (1 / \delta). \\ \end{array}
$$
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}:\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}\rightarrow \mathbb{R}^d,\mu_h^\star :\mathcal{X}\rightarrow \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}\to \mathbb{R}$ . The MDP is called low-rank if  $P_{h}$  is low-rank for all  $h\in [H]$ .
Suppose that we have a function class  $\Phi$  such that  $\phi_h^\star \in \Phi$  for all  $h$ , i.e.,  $\Phi$  is a realizable function class. For example, in linear MDPs, this is automatically satisfied since we know  $\phi^\star$  a priori, so  $\Phi$  is the singleton with  $\phi^\star$ . Having a realizable  $\Phi$  class is standard for solving low-rank MDPs [Uehara et al., 2021, Agarwal et al., 2023].
In what follows, let  $\alpha = \max_{h,\pi ,z,x,a}Z_h^\pi (z\mid x,a)$  denote the maximum density/mass value of the loss-to-go distributions. Note that  $\alpha \geq 1$  always since the mass at  $H + 1$  is deterministically placed at zero. If we further know that  $Z_{h}^{\pi}$  is discretely distributed, then  $\alpha = 1$ . If  $Z_{h}^{\pi}$  is continuously distributed, we assume it is bounded.
We consider the function class in Eq. (2), which we reproduce here:
$$
\mathcal {F} _ {h} ^ {\text {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 {3}
$$
$$
\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\}.
$$
The next lemma (Lemma D.1) shows that this function class satisfies distributional BC.
Lemma D.1.  $\mathcal{F}^{lin}$  satisfies distributional BC (Assumption 5.1).
Proof. We denote  $\| f\|_{\infty} = \max_{z,x,a}f(z|x,a)$ . For any  $f_{h + 1}\in \mathcal{F}_{h + 1}^{\mathrm{lin}}$ , we have  $\| f_{h + 1}\|_{\infty}\leq \alpha$  by the construction of  $\mathcal{F}_{h + 1}^{\mathrm{lin}}$ . Then, let  $\mathcal{T}^D$  be either the distributional Bellman operator or distributional optimality operator, the following equalities hold for the appropriate  $a^\prime (x^\prime)$  based on  $\mathcal{T}^D$
$$
\begin{array}{l} \mathcal {T} ^ {D} f _ {h + 1} (z \mid x, a) = \int_ {\mathcal {X}} \Pr_ {h} \left(x ^ {\prime} \mid x, a\right) \int_ {\mathbb {R}} \Pr_ {h} \left(c \mid x, a\right) f _ {h + 1} \left(z - c \mid x ^ {\prime}, a ^ {\prime} \left(x ^ {\prime}\right)\right) d x ^ {\prime} d c \\ = \left\langle \phi_ {h} ^ {\star} (x, a), \underbrace {\int_ {\mathcal {X}} \mu_ {h} \left(x ^ {\prime}\right) \int_ {\mathbb {R}} \operatorname * {P r} _ {h} \left(c \mid x , a\right) f _ {h + 1} \left(z - c \mid x ^ {\prime} , a ^ {\prime} \left(x ^ {\prime}\right)\right) \mathrm {d} c \mathrm {d} x ^ {\prime}} _ {:= w _ {h} (z)} \right\rangle \\ \end{array}
$$
Since  $\int_{\mathbb{R}}\operatorname *{Pr}_h(c\mid x,a)f_{h + 1}(z - c\mid x',a'(x'))\mathrm{d}c\leq \| f_{h + 1}\|_{\infty}$ , we know that
$$
\left\| w _ {h} (z) \right\| _ {2} \leq \left\| f _ {h + 1} \right\| _ {\infty} \sqrt {d} \leq \alpha \sqrt {d}.
$$
We further note that
$$
\max  _ {x, a, z} \left\langle \phi_ {h} ^ {\star} (x, a), w _ {h} (z) \right\rangle = \max  _ {x, a, z} \mathcal {T} ^ {D} f _ {h + 1} (z | x, a) \leq \| f _ {h + 1} \| _ {\infty} \leq \alpha .
$$
Also note that  $\phi_h^\star \in \Phi$  by realizability. Therefore,  $\mathcal{T}^D f_{h + 1}\in \mathcal{F}_h^{\mathrm{lin}}$ , which is the distributional BC condition.
D.1 Bounding the bracketing number via discretized rewards
We now bound the bracketing number of  $\mathcal{F}_h^{\mathrm{lin}}$  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  $\mathcal{F}_h^{\mathrm{lin}}$  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^\pi$  are discrete, we have  $\alpha = 1$ .
Now, let  $\varepsilon > 0$  be arbitrary and fixed. Recall that the  $\ell_{\infty}$  bracketing number is equivalent (up to universal constants) to the  $\ell_{\infty}$  covering number, so we will work with the latter. Let  $B(r)$  denote the  $d$ -dimensional ball of radius  $r$  (in  $\ell_2$ ). Recall that the  $\varepsilon$ -covering number (in  $\ell_2$ ) of functions  $[M] \mapsto B(r)$  scales as  $\mathcal{O}((r / \varepsilon)^{dM})$ . Let  $\mathcal{W}_{\varepsilon}$  be such the smallest cover. We can build a  $\ell_{\infty}$  cover of  $\mathcal{F}_h^{\mathrm{lin}}$  as follows:  $\mathcal{C}_{\varepsilon} = \{(x,a,z) \mapsto \langle \phi(x,a), w(z) \rangle, w \in \mathcal{W}_{\varepsilon}, \phi \in \Phi\}$ .
To check this is a  $\varepsilon$  cover, consider any  $f\in \mathcal{F}_h^{\mathrm{lin}}$ .  $f$  corresponds to some  $\phi$  and  $w$ . Let  $w^{\prime}$  be the neighbor of  $w$  in  $\mathcal{W}_{\varepsilon}$  and let  $f^{\prime}(x,a,z) = \langle \phi (x,a),w^{\prime}(z)\rangle$  so indeed  $f^{\prime}\in \mathcal{C}_{\varepsilon}$ . Then, for any  $x,a,z$ , we have  $|\langle \phi (x,a),w(z) - w'(z)\rangle |\leq \| \phi (x,a)\| _2\| w(z) - w'(z)\| _2\leq \varepsilon$ . Hence,  $\mathcal{C}_{\varepsilon}$  is an  $\ell_{\infty}$  cover of size  $\mathcal{O}((\sqrt{d} /\varepsilon)^{dM}\cdot |\Phi |)$ , and so we have shown that  $\log N_{\mathbb{I}}(\varepsilon ,\mathcal{F}_h^{\mathrm{lin}},\| \cdot \|_{\infty})\leq$ $\mathcal{O}(dM\log (d / \varepsilon) + \log |\Phi |)$ .
Linear MDPs: Recall that in linear MDPs, we know the true  $\phi^{\star}$  and so  $|\Phi| = 1$ . Thus, the bracketing number is simply  $\mathcal{O}(dM\log(d / \varepsilon))$  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  $\mathcal{X}$  be the context/feature space and  $\mathcal{V}$  be the label space, and we are given a dataset  $D = \{(x_{i},y_{i})\}_{i\in [n]}$  from a martingale process: for  $i = 1,2,\ldots ,n$  sample  $x_{i}\sim \mathcal{D}_{i}(x_{1:i - 1},y_{1:i - 1})$  and  $y_{i}\sim p(\cdot \mid x_{i})$ . Let  $f^{\star}(x,y) = p(y\mid x)$  and we are given a realizable, i.e.,  $f^{\star}\in \mathcal{F}$ , function class  $\mathcal{F}:\mathcal{X}\times \mathcal{V}\to \Delta (\mathbb{R})$  of distributions. The MLE is an estimate for  $f^{\star}$  that maximizes the log-likelihood objective over our dataset:
$$
\widehat {f} _ {\mathrm {M L E}} = \underset {f \in \mathcal {F}} {\arg \max } \sum_ {i = 1} ^ {n} \log f (x _ {i}, y _ {i}).
$$
For our guarantees to hold for general hypotheses classes  $\mathcal{F}$ , we use the bracketing number to quantify the statistical complexity of  $\mathcal{F}$  [van de Geer, 2000].