# Distributional Offline Policy Evaluation with Predictive Error Guarantees

Runzhe Wu111Email<EMAIL_ADDRESS>Masatoshi Uehara222Email:
<EMAIL_ADDRESS>Wen Sun333Email<EMAIL_ADDRESS>
(Cornell University)

###### Abstract

We study the problem of estimating the distribution of the return of a policy
using an offline dataset that is not generated from the policy, i.e.,
distributional offline policy evaluation (OPE). We propose an algorithm called
Fitted Likelihood Estimation (FLE), which conducts a sequence of Maximum
Likelihood Estimation (MLE) problems and has the flexibility of integrating
any state-of-art probabilistic generative models as long as it can be trained
via MLE. FLE can be used for both finite horizon and infinite horizon
discounted settings where rewards can be multi-dimensional vectors. In our
theoretical results, we show that for both finite and infinite horizon
discounted settings, FLE can learn distributions that are close to the ground
truth under total variation distance and Wasserstein distance, respectively.
Our theoretical results hold under the conditions that the offline data covers
the test policy’s traces and the supervised learning MLE procedures succeed.
Experimentally, we demonstrate the performance of FLE with two generative
models, Gaussian mixture models and diffusion models. For the multi-
dimensional reward setting, FLE with diffusion models is capable of estimating
the complicated distribution of the return of a test policy.

##  1 Introduction

Traditional Reinforcement Learning (RL) focuses on studying the expected
behaviors of a learning agent. However, modeling the expected behavior is not
enough for many interesting applications. For instance, when estimating the
value of a new medical treatment, instead of just predicting its expected
value, we may be interested in estimating the variance of the value as well.
For a self-driving car whose goal is to reach a destination as soon as
possible, in addition to predicting the expected traveling time, we may be
interested in estimating the tails of the distribution of traveling time so
that customers can prepare for worst-case situations. Other risk-sensitive
applications in finance and control often require one to model beyond the
expectation as well.

In this work, we study the question of how to estimate the distribution of the
return of a policy in Markov Decision Processes (MDPs) using only an offline
dataset that is not necessarily generated from the test policy (i.e.,
distributional offline policy evaluation). Estimating distributions of returns
has been studied in the setting called distributional RL [Bellemare et al.,
2017], where most of the existing works focus on solving the regular RL
problem, i.e., finding a policy that maximizes the expected return by treating
the task of predicting additional information beyond the mean as an auxiliary
task. Empirically, it is believed that this auxiliary task helps
representation learning which in turn leads to better empirical performance.
Instead of focusing on this auxiliary loss perspective, we aim to design
distributional OPE algorithms, which can accurately estimate the distribution
of returns with provable guarantees. We are also interested in the setting
where the one-step reward could be _multi-dimensional_ (i.e., multi-objective
RL), and the state/action spaces could be large or even continuous. This
requires us to design new algorithms that can leverage rich function
approximation (e.g., state-of-art probabilistic generative models).

Our algorithm, _Fitted Likelihood Estimation_ (FLE), is inspired by the
classic OPE algorithm Fitted Q Evaluation (FQE) [Munos and Szepesvári, 2008].
Given a test policy and an offline dataset, FLE iteratively calls a supervised
learning oracle — Maximum Likelihood Estimation (MLE) in this case, to fit a
conditional distribution to approximate a target distribution which is
constructed using the distribution learned from the previous iteration. At the
end of the training procedure, it outputs an estimator which aims to
approximate the true distribution of the return of the test policy. Our
algorithm is simple: like FQE, it decomposes the distributional OPE problem
into a sequence of supervised learning problems (in this case, MLE). Thus it
has great flexibility to leverage any state-of-art probabilistic generative
models as long as it can be trained via MLE. Such flexibility is important,
especially when we have large state/action space, and reward vectors coming
from complicated high-dimensional distributions. FLE naturally works for both
finite horizon setting and infinite horizon discounted setting.

Theoretically, we prove that our algorithm, FLE, can learn an accurate
estimator of the return distribution for both finite horizon MDPs and infinite
horizon discounted MDPs, under the assumptions that (1) _MLE can achieve good
in-distribution generalization bounds (i.e., supervised learning succeeds)_ ,
and (2) _the offline state-action distribution covers the test policy’s state-
action distribution_. The first condition is well studied in statistical
learning theory, and in practice, the state-of-the-art probabilistic
generative models trained via MLE (e.g., FLOW models [Dinh et al., 2014] and
Diffusion models [Sohl-Dickstein et al., 2015]) indeed also exhibit amazing
generalization ability. The second condition is necessary for offline RL and
is widely used in the regular offline RL literature (e.g., Munos and
Szepesvári [2008]). In other words, our analysis is modular: it simply
transfers the supervised learning MLE in-distribution generalization bounds to
a bound of distributional OPE. The accuracy of the estimator computed by FLE
is measured under total variation distance and p-Wasserstein distance, for
finite horizon setting and infinite horizon discounted setting, respectively.
To complete the picture, we further provide concrete examples showing that MLE
can provably have small in-distribution generalization errors in the analysis
section. To the best of our knowledge, this is the first PAC (Probably
Approximately Correct) learning algorithm for distributional OPE with general
function approximation.

Finally, we demonstrate our approach on a rich observation combination lock
MDP where it has a latent structure with the observations being high-
dimensional and continuous [Misra et al., 2020, Agarwal et al., 2020a, Zhang
et al., 2022]. We consider the setting where the reward comes from complicated
multi-dimensional continuous distributions (thus existing algorithms such as
quantile-regression TD [Dabney et al., 2018] do not directly apply here). We
demonstrate the flexibility of our approach by using two generative models in
FLE: the classic Gaussian mixture model and state-of-the-art diffusion model
[Ho et al., 2020].

###  1.1 Related Works

##### Distributional RL.

Quantile regression TD [Dabney et al., 2018] is one of the common approaches
for distributional OPE. A very recent work [Rowland et al., 2023] demonstrates
that quantile regression TD can converge to the TD fixed point solution of
which the existence is proved under an $\ell_{\infty}$-style norm (i.e.,
$\sup$ over all states). Rowland et al. [2023] do not consider the sample
complexity of OPE and the impact of learning from off-policy samples, and
their convergence analysis is asymptotic. Also, quantile regression TD only
works for scalar rewards. Another popular approach is categorical TD
[Bellemare et al., 2017], where one explicitly discretizes the return space.
However, for high-dimensional rewards, explicitly discretizing the return
space evenly can suffer the curse of dimensionality and fail to capture some
low-dimensional structures in the data distribution. Also, there is no
convergence or sample complexity analysis of the categorical algorithm for
OPE.

Ma et al. [2021] studied distributional offline policy optimization. In their
analysis, they focused on tabular MDPs with scalar rewards, and their
algorithm can learn a pessimistic estimate of the true inverse CDF of the
return. Keramati et al. [2020] also uses the distributional RL framework to
derive an approach that can optimistically estimate the CVaR value of a
policy’s return. Their analysis also only applies to tabular MDPs with scalar
rewards. In contrast, we focus on distributional OPE with general function
approximation beyond tabular or linear formats and MDPs with multi-dimensional
rewards.

Zhang et al. [2021] also consider learning from vector-valued rewards. They
propose a practical algorithm that minimizes the Maximum Mean Discrepancy
(MMD) without providing a sample complexity analysis. In contrast, we use MLE
to minimize total variation distance, and our error bound is based on total
variation distance. Note that a small total variation distance implies a small
MMD but not vice versa, which implies that our results are stronger.

##### Offline policy evaluation.

Fitted Q evaluation (FQE) [Munos and Szepesvári, 2008, Ernst et al., 2005] is
one of the most classic OPE algorithms. Many alternative approaches have been
recently proposed, such as minimax algorithms [Yang et al., 2020, Feng et al.,
2019, Uehara et al., 2020]. Somewhat surprisingly, algorithms based on FQE are
often robust and achieve stronger empirical performance in various benchmark
tasks [Fu et al., 2021, Chang et al., 2022]. Our proposed algorithm can be
understood as a direct generalization of FQE to the distributional setting.
Note sequential importance sampling approaches [Jiang and Li, 2016, Precup et
al., 2000] in regular RL have been applied to estimate distributions [Chandak
et al., 2021]. However, these methods suffer from the curse of the horizon,
i.e., the variance necessarily grows exponentially in the horizon.

##  2 Preliminaries

###  2.1 Finite-Horizon MDPs

We consider finite-horizon MDP with a vector-valued reward function, which is
a tuple $M(\mathcal{X},\mathcal{A},r,P,H,\mu)$ where $\mathcal{X}$ and
$\mathcal{A}$ are the state and action spaces, respectively, $P$ is the
transition kernel, $r$ is the reward function, i.e.,
$r(x,a)\in\Delta([0,1]^{d})$ where $d\in\mathbb{Z}^{+}$, $H$ is the length of
each episode, and $\mu\in\Delta(\mathcal{X})$ is the initial state
distribution. A policy is a mapping
$\pi:\mathcal{X}\rightarrow\Delta(\mathcal{A})$. We denote
$z\in\mathbb{R}^{d}$ as the accumulative reward vector across $H$ steps, i.e.,
$z=\sum_{h=1}^{H}r_{h}\in[0,H]^{d}$. Note that $z$ is a random vector whose
distribution is determined by a policy $\pi$ and the MDP. We denote
$Z^{\pi}\in\Delta([0,H]^{d})$ as the distribution444 Formally, they are called
probability density functions in the continuous setting and probability mass
functions in discrete settings, which are different from cumulative
distribution functions.  of the random variable $z$ under policy $\pi$. In
this paper, we are interested in estimating $Z^{\pi}$ using offline data. We
also define conditional distributions $Z^{\pi}_{h}(x,a)\in\Delta([0,H]^{d})$
which is the distribution of the return under policy $\pi$ starting with state
action $(x_{h},a_{h}):=(x,a)$ at time step $h$. It is easy to see that
$Z^{\pi}=\mathbb{E}_{x\sim\mu,a\sim\pi(x)}\left[Z^{\pi}_{1}(x,a)\right]$.

Given $\pi$, we also define $d^{\pi}_{h}$ as the state-action distribution
induced by policy $\pi$ at time step $h$, and
$d^{\pi}=\sum_{h=1}^{H}d^{\pi}_{h}/H$ as the average state-action distribution
induced by $\pi$.

We denote distributional Bellman operator [Morimura et al., 2012] associated
with $\pi$ as $\mathcal{T}^{\pi}$, which maps a conditional distribution to
another conditional distribution: given a state-action conditional
distribution $f\in\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,H]^{d})$, we
have
$\mathcal{T}^{\pi}f\in\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,H]^{d})$,
such that for any $(x,a,z)$:

$\displaystyle[\mathcal{T}^{\pi}f](z\,|\,x,a)=\mathbb{E}_{r\sim
r(s,a),x^{\prime}\sim
P(x,a),a^{\prime}\sim\pi(x)}\left[f\left(z-r|x^{\prime},a^{\prime}\right)\right].$

It is easy to see that for $\\{Z^{\pi}_{h}\\}_{h=1}^{H}$, we have
$\mathcal{T}^{\pi}Z^{\pi}_{h+1}=Z^{\pi}_{h}$ for all $h$.

###  2.2 Discounted Infinite-Horizon MDPs

The discounted infinite-horizon MDP is a tuple
$M(\mathcal{X},\mathcal{A},r,P,\gamma,\mu)$. The return vector is defined as
$z=\sum_{h=1}^{\infty}\gamma^{h-1}r_{h}$. We call $\gamma\in(0,1)$ the
discount factor. The distribution of return $z$ is thus
$Z^{\pi}\in\Delta([0,(1-\gamma)^{-1}]^{d})$. We also define the conditional
distributions $\bar{Z}^{\pi}(x,a)\in\Delta([0,(1-\gamma)^{-1}]^{d})$ which is
the distribution of the return under policy $\pi$ starting with state action
$(x,a)$. It is easy to see that
$Z^{\pi}=\mathbb{E}_{x\sim\mu,a\sim\pi(x)}\left[\bar{Z}^{\pi}(x,a)\right]$.

The state-action distribution of a given policy $\pi$ is also defined in a
discounted way:
$d^{\pi}=(1-\gamma)^{-1}\sum_{h=1}^{\infty}\gamma^{h-1}d^{\pi}_{h}$ where
$d^{\pi}_{h}$ is the state-action distribution induced by $\pi$ at time step
$h$.

The distributional Bellman operator maps a state-action conditional
distribution
$f\in\mathcal{X}\times\mathcal{A}\mapsto([0,(1-\gamma)^{-1}]^{d})$ to
$\mathcal{T}^{\pi}f\in\mathcal{X}\times\mathcal{A}\mapsto([0,(1-\gamma)^{-1}]^{d})$,
such that for any $(x,a,z)$:

$\displaystyle[\mathcal{T}^{\pi}f](z\,|\,x,a)=\mathbb{E}_{r\sim
r(s,a),x^{\prime}\sim
P(x,a),a^{\prime}\sim\pi(x)}\left[f\left(\frac{z-r}{\gamma}\,\middle|\,x^{\prime},a^{\prime}\right)\right].$

We can verify that $\bar{Z}^{\pi}$ is a fixed point of the distributional
Bellman operator, i.e., $\mathcal{T}^{\pi}\bar{Z}^{\pi}=\bar{Z}^{\pi}$.

###  2.3 Offline Policy Evaluation Setup

We consider estimating the distribution $Z^{\pi}$ using offline data which
does not come from $\pi$ (i.e., off-policy setting). We assume we have a
dataset $\mathcal{D}=\\{x_{i},a_{i},r_{i},x^{\prime}_{i}\\}_{i=1}^{n}$ that
contains i.i.d. tuples, such that
$x,a\sim\rho\in\Delta(\mathcal{X}\times\mathcal{A})$, $s^{\prime}\sim
P(\cdot|s,a)$, and $r\sim r(s,a)$. For finite-horizon MDPs, we randomly and
evenly split $\mathcal{D}$ into $H$ subsets,
$\mathcal{D}_{1},\dots,\mathcal{D}_{H}$, for the convenience of analysis. Each
subset contains $n/H$ samples. For infinite-horizon MDPs, we split it into $T$
subsets in the same way. Here $T$ is the number of iterations which we will
define later.

We consider learning distribution $Z^{\pi}$ via general function
approximation. For finite-horizon MDPs, we denote $\mathcal{F}_{h}$ as a
function class that contains state-action conditional distributions, i.e.,
$\mathcal{F}_{h}\subset\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,H]^{d})$,
which will be used to learn $Z^{\pi}_{h}$. For infinite-horizon MDPs, we
assume a function class
$\mathcal{F}\subset\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,(1-\gamma)^{-1}]^{d})$.

##### Notations.

Given distributions $P_{1}$ and $P_{2}$ on a set $\mathcal{S}$, we denote
$d_{tv}$ as the total variation distance between the two distributions, i.e.,
$d_{tv}(P_{1},P_{2})=\|P_{1}-P_{2}\|_{1}/2$. We denote $d_{w,p}$ as the
$p$-Wasserstein distance for which the metric is induced by $\ell_{2}$ norm,
i.e.,
$d_{w,p}(P_{1},P_{2})=(\inf_{c\in\mathcal{C}}\mathop{\mathbb{E}}_{x,y\sim
c}\|x-y\|^{p})^{1/p}$ where $\mathcal{C}$ denotes the set of all couplings of
$P_{1}$ and $P_{2}$. We note that $d_{tv}$ dominates $d_{w,p}$ when the
support is bounded (see Lemma C.6 for details):

$d_{w,p}^{p}(P_{1},P_{2})\leq\text{\rm diam}^{p}(\mathcal{S})\cdot
d_{tv}(P_{1},P_{2}).$ (1)

where $\text{\rm diam}(\mathcal{S})=\sup_{x,y\in\mathcal{S}}\|x-y\|$ is the
diameter of $\mathcal{S}$.

##  3 Fitted Likelihood Estimation

In this section, we present our algorithm — _Fitted Likelihood Estimation_
(FLE) for distributional OPE. Algorithm 1 is for finite-horizon MDPs and
Algorithm 2 is for infinite-horizon MDPs. We first introduce the former.

Algorithm 1 takes the offline dataset
$\mathcal{D}=\\{\mathcal{D}_{h}\\}_{h=1}^{H}$ and the function class
$\\{\mathcal{F}_{h}\\}_{h=1}^{H}$ as inputs and iteratively performs Maximum
likelihood estimation (MLE) starting from $H$ to time step $h=1$. Consider a
particular time step $h$, given $\hat{f}_{h+1}$ which is learned from the
previous iteration, FLE treats $\mathcal{T}^{\pi}\hat{f}_{h+1}$ as the target
distribution to fit. To learn $\mathcal{T}^{\pi}\hat{f}_{h+1}$, it first
generates samples from it (Line 6), which is doable as long as we can generate
samples from the conditional distribution $\hat{f}_{h+1}(\cdot|x,a)$ given any
$(x,a)$. Once we generate samples from $\mathcal{T}^{\pi}\hat{f}_{h+1}$, we
fit $\hat{f}_{h}$ to estimate $\mathcal{T}^{\pi}\hat{f}_{h+1}$ by MLE (Line
13). The algorithm returns $\hat{f}_{1}$ to approximate $Z^{\pi}_{1}$. To
estimate $Z^{\pi}$, we can compute
$\mathbb{E}_{x\sim\mu,a\sim\pi(x)}\hat{f}_{1}(x,a)$, recalling $\mu$ is the
initial state distribution.

Algorithm 2 is quite similar to Algorithm 1 but is for infinite-horizon MDPs,
and it has two distinctions. First, we introduce the discount factor $\gamma$.
Second, compared to Algorithm 1 where we perform MLE in a backward manner
(from $h=H$ to $1$), for infinite-horizon MDPs, we repeatedly apply MLE in a
time-independent way. Particularly, it treats $\mathcal{T}^{\pi}\hat{f}_{t-1}$
as the target distribution to fit by MLE at round $t$. To finally estimate
$Z^{\pi}$, we can compute $\mathbb{E}_{x\sim\mu,a\sim\pi(x)}\hat{f}_{T}(x,a)$.

To implement either algorithm, we need function $f$ that has the following two
properties: (1) it can generate samples given any state-action pair, i.e.,
$z\sim f(\cdot|x,a)$, and (2) given any triple $(x,a,z)$ we can evaluate the
conditional likelihood, i.e., we can compute $f(z|x,a)$. Such function
approximation $f$ is widely available in practice, including discrete
histogram-based models, Gaussian mixture models, Flow models [Dinh et al.,
2014], and diffusion model [Sohl-Dickstein et al., 2015]. Indeed in our
experiment, we implement FLE with Gaussian mixture models and diffusion models
[Ho et al., 2020], both of which are optimized via MLE.

Algorithm 1 Fitted Likelihood Estimation (FLE) for finite-horizon MDPs

1: Input: dataset $\\{\mathcal{D}_{h}\\}_{h=1}^{H}$ and function classes
$\\{\mathcal{F}_{h}\\}_{h=1}^{H}$

2: for $h=H,H-1,\dots,1$ do

3: $\mathcal{D}_{h}^{\prime}=\emptyset$

4: for $x,a,r,x^{\prime}\in\mathcal{D}_{h}$ do

5: if $h<H$ then

6: $a^{\prime}\sim\pi(x^{\prime})$,
$y\sim\hat{f}_{h+1}(\cdot\,|\,x^{\prime},a^{\prime})$

7: Set $z=r+y$

8: else

9: Set $z=r$

10: end if

11: $\mathcal{D}^{\prime}_{h}=\mathcal{D}^{\prime}_{h}\cup\\{(x,a,z)\\}$

12: end for

13:
$\hat{f}_{h}=\mathop{\arg\max}_{f\in\mathcal{F}_{h}}\sum_{(x,a,z)\in\mathcal{D}^{\prime}_{h}}\log
f(z\,|\,x,a)$

14: end for

Algorithm 2 Fitted Likelihood Estimation (FLE) for infinite-horizon MDPs

1: Input: dataset $\\{\mathcal{D}_{t}\\}_{t=1}^{T}$ and function classes
$\mathcal{F}$

2: for $t=1,2,\dots,T$ do

3: $\mathcal{D}_{t}^{\prime}=\emptyset$

4: for $x,a,r,x^{\prime}\in\mathcal{D}_{t}$ do

5: $a^{\prime}\sim\pi(x^{\prime})$

6: $y\sim\hat{f}_{t-1}(\cdot\,|\,x^{\prime},a^{\prime})$

7: $z=r+\gamma y$

8: $\mathcal{D}^{\prime}_{t}=\mathcal{D}^{\prime}_{t}\cup\\{(x,a,z)\\}$

9: end for

10:
$\hat{f}_{t}=\mathop{\arg\max}_{f\in\mathcal{F}}\sum_{(x,a,z)\in\mathcal{D}^{\prime}_{t}}\log
f(z\,|\,x,a)$

11: end for

Regarding computation, the main bottleneck is the MLE step (Line 13 and 10).
While we present it with a $\mathop{\arg\max}$ oracle, in both practice and
theory, an approximation optimization oracle is enough. In theory, as we will
demonstrate, as long as we can find some $\hat{f}_{h}$ that exhibits good in-
distribution generalization bound (i.e.,
$\mathbb{E}_{x,a\sim\rho}d_{tv}(\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a))$
or
$\mathbb{E}_{x,a\sim\rho}d_{tv}(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a))$
is small), then we can guarantee to have an accurate estimator for $Z^{\pi}$.
Note that here $\rho$ is the training distribution for MLE, thus we care about
in-distribution generalization. Thus our approach is truly a reduction to
supervised learning: as long as the supervised learning procedure (in this
case, MLE) learns a model with good in-distribution generalization
performance, we can guarantee good prediction performance for FLE. Any
advancements from training generative models via MLE (e.g., better training
heuristics and better models) thus can immediately lead to improvement in
distributional OPE.

######  Remark 3.1 (Comparison to prior models).

The categorical algorithm [Bellemare et al., 2017] works by minimizing the
cross-entropy loss between the (projected) target distribution and the
parametric distribution, which is equivalent to maximizing the likelihood of
the parametric model.

######  Remark 3.2 (FQE as a special instance).

When reward is only a scalar, and we use fixed-variance Gaussian distribution
$f(\cdot|x,a):=\mathcal{N}(g(x,a),\sigma^{2})$ where
$g:\mathcal{X}\times\mathcal{A}\mapsto[0,H]$, and $\sigma>0$ is a fixed (not
learnable) parameter, MLE becomes a least square oracle, and FLE reduces to
FQE — the classic offline policy evaluation algorithm.

##  4 Theoretical Analysis

In this section, we present the theoretical guarantees of FLE. As a warmup, we
start by analyzing the performance of FLE for the finite horizon setting
(Section 4.1), where we bound the prediction error using total variation
distance. Then we study the guarantees of FLE for the infinite-horizon
discounted scenario in Section 4.2 where the prediction error is measured
under $p$-Wasserstein distance. Note that from Equation 1, TD distance
dominates $p$-Wasserstein distance which indicates that our guarantee for the
finite horizon setting is stronger. This shows an interesting difference
between the two settings. Finally, we present some concrete examples in
Section 4.3. All proofs can be found in Appendix A.

###  4.1 Finite Horizon

We start by stating the key assumption for OPE, which concerns the overlap
between $\pi$’s distribution and the offline distribution $\rho$.

######  Assumption 4.1 (Coverage).

We assume there exists a constant $C$ such that for all $h\in[H]$ the
following holds

$\displaystyle\sup_{\begin{subarray}{c}f_{h}\in\mathcal{F}_{h}\\\
f_{h+1}\in\mathcal{F}_{h+1}\end{subarray}}\frac{\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{h}}d^{2}_{tv}\left(f_{h}(x,a),[\mathcal{T}^{\pi}f_{h+1}](x,a)\right)}{\mathop{\mathbb{E}}_{x,a\sim\rho}d^{2}_{tv}\left(f_{h}(x,a),[\mathcal{T}^{\pi}f_{h+1}](x,a)\right)}\leq
C.$

The data coverage assumption is necessary for off-policy learning. Assumption
4.1 incorporates the function class into the definition of data coverage and
is always no larger than the usual density ratio-based coverage definition,
i.e., $\sup_{h,x,a}d_{h}^{\pi}(x,a)/\rho(x,a)$ which is a classic coverage
measure in offline RL literature (e.g., Munos and Szepesvári [2008]). This
type of refined coverage is used in the regular RL setting [Xie et al., 2021,
Uehara and Sun, 2021].

Next, we present the theoretical guarantee of our approach under the
_assumption that MLE can achieve good supervised learning style in-
distribution generalization bound_. Recall that in each iteration of our
algorithm, we perform MLE to learn a function $\hat{f}_{h}$ to approximate the
target $\mathcal{T}^{\pi}\hat{f}_{h+1}$ under the training data from $\rho$.
By supervised learning style in-distribution generalization error, we mean the
divergence $d_{tv}$ between $\hat{f}_{h}$ and the target
$\mathcal{T}^{\pi}\hat{f}_{h+1}$ under the _training distribution $\rho$_.
Such an in-distribution generalization bound for MLE is widely studied in
statistical learning theory literature [Van de Geer, 2000, Zhang, 2006], and
used in RL literature (e.g., Agarwal et al. [2020b], Uehara et al. [2021],
Zhan et al. [2022]). The following theorem demonstrates a reduction framework:
as long as supervised learning MLE works, our estimator of $Z^{\pi}$ is
accurate.

######  Theorem 4.2.

Under 4.1, suppose we have a sequence of functions
$\hat{f}_{1},\dots,\hat{f}_{H}:\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,H]^{d})$
and a sequence of values $\zeta_{1},\dots,\zeta_{H}\in\mathbb{R}$ such that

$\textstyle\bigg{(}\mathop{\mathbb{E}}_{x,a\sim\rho}\
d^{2}_{tv}\left(\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a)\right)\bigg{)}^{1/2}\leq\zeta_{h}$

holds for all $h\in[H]$. Let our estimator
$\hat{f}\coloneqq\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\hat{f}_{1}(x,a)$.
Then we have

$\displaystyle
d_{tv}\left(\hat{f},Z^{\pi}\right)\leq\sqrt{C}\sum_{h=1}^{H}\zeta_{h}.$

Here recall that $C$ is the coverage definition. Thus the above theorem
demonstrates that when $\rho$ covers $d^{\pi}$ (i.e., $C<\infty$), small
supervised learning errors (i.e., $\zeta_{h}$) imply small prediction error
for distributional OPE.

Now to complete the picture, we provide some sufficient conditions where MLE
can achieve small in-distribution generalization errors. The first condition
is stated below.

######  Assumption 4.3 (Bellman completeness).

We assume the following holds:

$\max_{h\in[H],f\in\mathcal{F}_{h+1}}\min_{g\in\mathcal{F}_{h}}\mathbb{E}_{x,a\sim\rho}\
d_{tv}\big{(}g(x,a),[\mathcal{T}^{\pi}f](x,a)\big{)}=0.$

We call the LHS of the above inequality _inherent (distributional) Bellman
error_.

This condition ensures that in each call of MLE in our algorithm, our function
class $\mathcal{F}_{h}$ contains the target $\mathcal{T}^{\pi}\hat{f}_{h+1}$.
It is possible to relax the above condition to a setting where the inherent
Bellman error is bounded by some small number $\delta$ (i.e., for MLE, this
corresponds to agnostic learning where the hypothesis class may not contain
the target, which is also a well-studied problem in statistical learning
theory [Van de Geer, 2000]). Here we mainly focus on $\delta=0$ case.

The Bellman completeness assumption (or, more generally, inherent Bellman
error being small) is standard in offline RL literature [Munos and Szepesvári,
2008]. Indeed, in the regular RL setting, when learning with off-policy data,
without such a Bellman completeness condition, algorithms such as TD learning
or value iteration-based approaches (e.g., FQE) can diverge [Tsitsiklis and
Van Roy, 1996], and TD fixed solution can be arbitrarily bad in terms of
approximating the true value (e.g., Munos [2003], Scherrer [2010], Kolter
[2011]). Since distributional RL generalizes regular RL, to prove convergence
and provide an explicit sample complexity, we also need such a Bellman
completeness condition.

The second condition is the bounded complexity of $\mathcal{F}_{h}$. A simple
case is when $\mathcal{F}$ is discrete where the standard statistical
complexity of $\mathcal{F}$ is $\ln(|\mathcal{F}_{h}|)$. We show the following
result for MLE’s in-distribution generalization error.

######  Lemma 4.4.

Assume $|\mathcal{F}_{h}|<\infty$. For FLE (Algorithm 1), under 4.3, MLEs have
the following guarantee:

$\displaystyle\mathop{\mathbb{E}}_{x,a\sim\rho}d_{tv}^{2}\Big{(}\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}]$
$\displaystyle(x,a)\Big{)}\leq\frac{4H}{n}\log(|\mathcal{F}_{h}|H/\delta)$

for all $h\in[H]$ with probability at least $1-\delta$.

For infinite hypothesis classes, we use bracketing number [Van de Geer, 2000]
to quantify the statistical complexities.

######  Definition 4.5 (Bracketing number).

Consider a function class $\mathcal{F}$ that maps $\mathcal{X}$ to
$\mathbb{R}$. Given two functions $l$ and $u$, the bracket $[l,u]$ is the set
of all functions $f\in\mathcal{F}$ with $l(x)\leq f(x)\leq u(x)$ for all
$x\in\mathcal{X}$. An $\epsilon$-bracket is a bracket $[l,u]$ with
$\|l-u\|\leq\epsilon$. The bracketing number of $\mathcal{F}$ w.r.t. the
metric $\|\cdot\|$ denoted by $N_{[]}(\epsilon,\mathcal{F},\|\cdot\|)$ is the
minimum number of $\epsilon$-brackets needed to cover $\mathcal{F}$.

We can bound MLE’s generalization error using bracket number of $\mathcal{F}$.

######  Lemma 4.6.

For FLE (Algorithm 1), under Assumption 4.3, we have

$\displaystyle\mathop{\mathbb{E}}_{x,a\sim\rho}d_{tv}^{2}\Big{(}\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a)\Big{)}\leq\frac{10H}{n}\log\left(N_{[]}\big{(}(nH^{d})^{-1},\mathcal{F}_{h},\|\cdot\|_{\infty}\big{)}H/\delta\right)$

for all $h\in[H]$ with probability at least $1-\delta$.

With the generalization bounds of MLE, via Theorem 4.2, we can derive the
following specific error bound for FLE.

######  Corollary 4.7.

Under Assumption 4.1 and 4.3, for FLE (Algorithm 1), with probability at least
$1-\delta$, we have

$\displaystyle
d_{tv}\left(\hat{f},Z^{\pi}\right)\leq\sqrt{C}\sum_{h=1}^{H}\sqrt{\frac{4H}{n}\log(|\mathcal{F}_{h}|H/\delta)}$

when $|\mathcal{F}_{h}|<\infty$ for all $h\in[H]$, and

$\displaystyle
d_{tv}\Big{(}\hat{f},Z^{\pi}\Big{)}\leq\sqrt{C}\sum_{h=1}^{H}\sqrt{\frac{10H}{n}\log\left(N_{[]}\big{(}(nH^{d})^{-1},\mathcal{F}_{h},\|\cdot\|_{\infty}\big{)}H/\delta\right)}.$

for infinite function class $\mathcal{F}_{h}$.

Overall, our theory indicates that if we can train accurate distributions
(e.g., generative models) via supervised learning (i.e., MLE here), we
automatically have good predictive performance on estimating $Z^{\pi}$. This
provides great flexibility for designing special algorithms.

######  Remark 4.8 (Offline CVaR Estimation).

As a simple application, FLE can derive an estimator for the CVaR of the
return under the test policy $\pi$. This is doable because CVaR is Lipschitz
with respect to distributions in total variation distance, and thus our
results can be directly transferred. See Appendix B for details. Essentially,
any quantity that is Lipschitz with respect to distributions in total
variation distance can be estimated using our method and the error bound of
FLE directly applies.

###  4.2 Infinite Horizon

Next we introduce the theoretical guarantees of FLE for infinite horizon MDPs.
Although the idea is similar, there is an obstacle: we can no longer obtain
guarantees in terms of the total variation distance. This is perhaps not
surprising considering that the distributional Bellman operator for discounted
setting is not contractive in total variation distance [Bellemare et al.,
2017]. Fortunately, we found the Bellman operator is contractive under the
Wasserstein distance measure. Note that the contractive result we established
under Wasserstein distance is different from previous works [Bellemare et al.,
2017, 2023, Zhang et al., 2021] in that these previous works consider the
supremum Wasserstein distance: $\sup_{x,a}d_{w,p}$, while our contrative
property is measured under an average Wasserstein distance:
$(\mathop{\mathbb{E}}_{x,a\sim d^{\pi}}d^{2p}_{w,p})^{1/(2p)}$ which is
critical to get a sample complexity bound for distributional OPE. More
formally, the following lemma summarizes the contractive property.

######  Lemma 4.9.

The distributional Bellman operator is $\gamma^{1-1/(2p)}$-contractive under
the metric $(\mathop{\mathbb{E}}_{x,a\sim d^{\pi}}d^{2p}_{w,p})^{1/(2p)}$,
i.e., for any
$f,f^{\prime}\in\mathcal{X}\times\mathcal{A}\mapsto[0,(1-\gamma)^{-1}]^{d}$,
it holds that

$\displaystyle\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left([\mathcal{T}^{\pi}f](x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)\right)^{\frac{1}{2p}}\leq\gamma^{1-\frac{1}{2p}}\cdot\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(f(x,a),f^{\prime}(x,a)\right)\right)^{\frac{1}{2p}}.$

We note that the contractive result in $\sup_{s,a}d_{w,p}$ does not imply the
result in the above lemma, thus cannot be directly applied to the OPE setting.

Additionally, since the total variation distance dominates the Wasserstein
distance on bounded sets (see (1)), the estimation error of MLE measured under
total variation distance can be used to bound the Wasserstein distance. This
allows us to derive theoretical guarantees for FLE on Wasserstein distance. To
this end, we start again with the coverage assumption, which is similar to
4.1. Note that we have replaced the total variation distance with the
Wasserstein distance.

######  Assumption 4.10 (Coverage).

We assume there exists a constant $C$ such that the following holds

$\displaystyle\sup_{\begin{subarray}{c}f,f^{\prime}\in\mathcal{F}\end{subarray}}\frac{\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(f(x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)}{\mathop{\mathbb{E}}_{x,a\sim\rho}d^{2p}_{w,p}\left(f(x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)}\leq
C.$

As similar to Theorem 4.2, the following theorem states that as long as the
supervised learning is accurate, our estimator of $Z^{\pi}$ will be accurate
as well, but under $p$-Wasserstein distance this time.

######  Theorem 4.11.

Under 4.10, suppose we have a sequence of functions
$\hat{f}_{1},\dots,\hat{f}_{T}:\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,(1-\gamma)^{-1}]^{d})$
and an upper bound $\zeta\in\mathbb{R}$ such that

$\textstyle\bigg{(}\mathop{\mathbb{E}}_{x,a\sim\rho}\
d_{w,p}^{2p}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right)\bigg{)}^{\frac{1}{2p}}\leq\zeta$

holds for all $t\in[T]$. Let our estimator
$\hat{f}\coloneqq\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\hat{f}_{T}(x,a)$.
Then we have, for all $p\geq 1$,

$d_{w,p}\left(\hat{f},Z^{\pi}\right)\leq\frac{2C^{\frac{1}{2p}}}{(1-\gamma)^{\frac{3}{2}}}\cdot\zeta+\frac{\sqrt{d}\cdot\gamma^{\frac{T}{2}}}{(1-\gamma)^{\frac{3}{2}}}.$
(2)

The upper bound in (2) is actually a simplified version as we aim to present a
cleaner result. For a more refined upper bound that has more $p$-dependent
terms, please refer to Theorem A.2 in the appendix.

For the first additive term in $\eqref{eq:geo-zeta_t}$, we will later
demonstrate that the $\zeta$ derived from MLE depends on $p^{-1}$ at an
exponential rate. The second term is relatively insignificant as it converges
to zero at the rate of $\gamma^{T/2}$.

To proceed, we introduce the Bellman completeness assumption for infinite-
horizon MDPs, which is one of the sufficient conditions for MLE to achieve
small in-distribution generalization errors.

######  Assumption 4.12 (Bellman completeness).

We assume the following holds:

$\max_{f\in\mathcal{F}}\min_{g\in\mathcal{F}}\mathbb{E}_{x,a\sim\rho}\
d_{w,p}\big{(}g(x,a),[\mathcal{T}^{\pi}f](x,a)\big{)}=0.$

As similar to the previous section, when Bellman completeness holds and the
function class has bounded complexity, MLE can get a small generalization
error, which is stated in the following lemma.

######  Lemma 4.13.

For FLE (Algorithm 2), under 4.12, by applying MLEs we have, for all
$t\in[T]$,

$\displaystyle\mathop{\mathbb{E}}_{x,a\sim\rho}d_{w,p}^{2p}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right)\leq\left(\frac{\sqrt{d}}{1-\gamma}\right)^{2p}\frac{4T}{n}\log(|\mathcal{F}|T/\delta)$

when $|\mathcal{F}|<\infty$, and

$\displaystyle\mathop{\mathbb{E}}_{x,a\sim\rho}d_{w,p}^{2p}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right)\leq\left(\frac{\sqrt{d}}{1-\gamma}\right)^{2p}\frac{10T}{n}\log\left(N_{[]}\left(\frac{(1-\gamma)^{d}}{n},\mathcal{F},\|\cdot\|_{\infty}\right)T\Big{/}\delta\right)$

when $|\mathcal{F}|=\infty$, with probability at least $1-\delta$.

The multiplicative term $T$ in the upper bounds above comes from the data
splitting (recall that we have split the dataset $\mathcal{D}$ into $T$
subsets: $\mathcal{D}_{1},\dots,\mathcal{D}_{T}$). We think that a more
careful analysis may be able to remove the need of data splitting and thus
eliminate the term $T$ here, leading to a slightly better polynomial
dependence on the effective horizon $1/(1-\gamma)$ in the final sample
complexity bound. We leave this for future work.

In view of the above mentioned, to derive the specific error bound of FLE, we
need to choose an appropriate $T$ to make a good balance. The $T$ we choose is
of the logarithmic order. The result is shown in the corollary below.

######  Corollary 4.14.

We define

$\displaystyle\iota=\begin{cases}\log(|\mathcal{F}|/\delta),&\text{if\quad}|\mathcal{F}|<\infty;\\\
\log\left(N_{[]}\left(\frac{(1-\gamma)^{d}}{n},\mathcal{F},\|\cdot\|_{\infty}\right)\Big{/}\delta\right),&\text{if\quad}|\mathcal{F}|=\infty.\end{cases}$

Then under Assumption 4.10 and 4.12, for FLE (Algorithm 2), if we pick

$T=\log\left(C^{\frac{1}{2p}}\cdot\iota^{\frac{1}{2p}}\cdot\left(1-\gamma^{\frac{1}{2}}\right)^{-1}\cdot
n^{-\frac{1}{2p}}\right)\Big{/}\log\left(\gamma^{1-\frac{1}{2p}}\right)$

then with probability at least $1-\delta$, we have

$\displaystyle
d_{w,p}\left(\hat{f},Z^{\pi}\right)\leq\widetilde{O}\left(\frac{C^{\frac{1}{2p}}\cdot\iota^{\frac{1}{2p}}\cdot\sqrt{d}}{(1-\gamma)^{\frac{5}{2}}}\cdot
n^{-\frac{1}{2p}}\right)$

where
$\hat{f}\coloneqq\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\hat{f}_{T}(x,a)$.

The above upper bound depends on $n^{-1/(2p)}$, which seems unsatisfactory,
especially when $p$ is large. However, we believe that it is actually tight
since the previous study has shown that the minimax rate of estimating
$d_{w,p}$ using i.i.d samples from the given distribution is around
$O(n^{-1/(2p)})$ [Singh and Póczos, 2018]. More formally, given a distribution
$Q$ and $n$ i.i.d samples from $Q$, any algorithm that maps the $n$ i.i.d
samples to a distribution $\hat{Q}$, must have
$d_{w,p}(\hat{Q},Q)=\widetilde{\Omega}(n^{-1/(2p)})$ in the worst case. Note
that distributional OPE is strictly harder than this problem.

###  4.3 Examples

In this section, we discuss two examples: one is tabular MDPs, and the other
one is Linear Quadratic Regulators. For simplicity of presentation, we focus
on scalar rewards and finite horizon.

####  4.3.1 Tabular MDPs

The example we consider is tabular MDP (i.e., $|\mathcal{X}|$ and
$|\mathcal{A}|$ are finite) with continuous known reward distributions.
Specifically, we consider the sparse reward case where we only have a reward
at the last time step $H$ and have zero rewards at time step $h<H$. For each
$(x,a)$, Denote $r_{H}(x,a)\in\Delta([0,1])$.

Note that in this setup, via induction, it is easy to verify that for any
$h,x,a$, $Z^{\pi}_{h}(\cdot|x,a)$ is a mixture of the distributions
$\\{r_{H}(x,a):x\in\mathcal{X},a\in\mathcal{A}\\}$, i.e., for any $h,x,a$,
there exists a probability weight vector
$w\in\Delta(|\mathcal{X}||\mathcal{A}|)$, such that
$Z^{\pi}_{h}(\cdot|x,a)=\sum_{x^{\prime},a^{\prime}\in\mathcal{X}\times\mathcal{A}}w(x^{\prime},a^{\prime})r_{H}(\cdot|x^{\prime},a^{\prime})$.
Note that the parameters $w(x,a)$ are unknown due to the unknown transition
operator $P$, and need to be learned. Thus, in this case, we can design
function class $\mathcal{F}_{h}$ as follows:

$\displaystyle\mathcal{F}_{h}=\bigg{\\{}f(\cdot|x,a)=\sum_{x^{\prime},a^{\prime}\in\mathcal{X}\times\mathcal{A}}w_{x,a}(x^{\prime},a^{\prime})r_{H}(\cdot|x^{\prime},a^{\prime}):\big{\\{}w_{x,a}\in\Delta(|\mathcal{X}||\mathcal{A}|)\big{\\}}_{x,a\in\mathcal{X}\times\mathcal{A}}\bigg{\\}}.$

It is not hard to verify that $\\{\mathcal{F}_{h}\\}_{h=1}^{H}$ does satisfy
the Bellman complete condition. The log of the bracket number of
$\mathcal{F}_{h}$ is polynomial with respect to $|\mathcal{X}||\mathcal{A}|$.

######  Lemma 4.15.

In the above example, the complexity of $\mathcal{F}_{h}$ in bounded: $\log
N_{[]}(\epsilon,\mathcal{F}_{h},\|\cdot\|_{\infty})\leq
O(|\mathcal{X}|^{2}|\mathcal{A}|^{2}\log(r_{\infty}|\mathcal{X}||\mathcal{A}|/\epsilon))$
where $r_{\infty}\coloneqq\|r_{H}\|_{\infty}$.

Thus Algorithm 1 is capable of finding an accurate estimator of $Z^{\pi}$ with
sample complexity scaling polynomially with respect to the size of the state
and action spaces and horizon.

####  4.3.2 Linear Quadratic Regulator

The second example is LQR. We have
$\mathcal{X}\subset\mathbb{R}^{d_{x}},\mathcal{A}\subset\mathbb{R}^{d_{a}}$.

$\displaystyle x_{h+1}=Ax_{h}+Ba_{h},\quad
r(x_{h},a_{h})=-(x_{h}^{\top}Qx_{h}+a_{h}^{\top}Ra_{h})+\varepsilon$

where $\varepsilon\sim\mathcal{N}(0,\sigma^{2})$. Since the optimal policy for
LQR is a linear policy, we consider evaluating a linear policy $\pi(x):=Kx$
where $K\in\mathbb{R}^{d_{a}\times d_{x}}$. For this linear policy,
$Z^{\pi}_{h}(\cdot|x,a)$ is a Gaussian distribution, i.e.,
$Z^{\pi}_{h}(\cdot|x,a)=\mathcal{N}(\mu_{h}(x,a),\sigma_{h}(x,a))$, where
$\mu_{h}(x,a)$ and $\sigma_{h}(x,a)$ has closed form solutions.

######  Lemma 4.16.

For LQR defined above, $\mu_{h}(x,a)$ and $\sigma_{h}(x,a)$ has the following
closed form solutions

$\displaystyle\mu_{h}(x,a)=-(Ax+Ba)^{\top}U_{h+1}(Ax+Ba)-x^{\top}Qx-a^{\top}Ra,\quad\sigma^{2}_{h}(x,a)=(H-h+1)\sigma^{2}$

where we denote
$U_{h}=\sum_{i=h}^{H}((A+BK)^{i-h-1})^{\top}(Q+K^{\top}RK)(A+BK)^{i-h-1}$.

Thus our function class $\mathcal{F}_{h}$ can be designed as follows:

$\displaystyle\mathcal{F}_{h}=\Big{\\{}f(\cdot|x,a)=\mathcal{N}\big{(}\cdot\,\big{|}\,x^{\top}M_{1}x+a^{\top}M_{2}x+a^{\top}M_{3}a,(H-h+1)\sigma^{2}\big{)},\
\forall M_{1},M_{2},M_{3}\Big{\\}}$

We can show that this function class satisfies Bellman completeness.
Furthermore, here, we can refine $C$ in Assumption 4.1 to a relative condition
number following the derivation in Uehara and Sun [2021]. More specifically,
$C$ is $\sup_{w\neq
0,h}\frac{w^{\top}\mathbb{E}_{d^{\pi}_{h}}[\phi(x,a)\phi^{\top}(x,a)]w}{w^{\top}\mathbb{E}_{\rho}[\phi(x,a)\phi^{\top}(x,a)]w}$
where $\phi(x,a)=(x^{\top},a^{\top})^{\top}\otimes(x^{\top},a^{\top})^{\top}$
is a quadratic feature and $\otimes$ is the Kronecker product. Under some
regularity assumption (i.e., norm $M_{1},M_{2},M_{3}$ are bounded, which is
the case when the dynamical system induced by the linear policy is stable),
this function class has bounded statistical complexity.

######  Lemma 4.17.

We assume there exist parameters $m_{x},m_{a},m_{1},m_{2},m_{3}$ for which
$\|x\|_{2}\leq m_{x}$ for all $x\in\mathcal{X}$ and $\|a\|_{2}\leq m_{a}$ for
all $a\in\mathcal{A}$, and $\|M_{i}\|_{\rm F}\leq m_{i}$ for $i=1,2,3$. Then
we have

$\log N_{[]}(\epsilon,\mathcal{F}_{h},\|\cdot\|_{\infty})\leq{\rm
Poly}\left(d_{x},d_{a},\log\frac{m_{x}m_{a}m_{1}m_{2}m_{3}}{\epsilon\sigma}\right).$

It is unclear if quantile regression TD or categorical TD can achieve
meaningful guarantees on LQR since the function classes used by them do not
satisfy Bellman completeness (i.e., given any conditional density
$f(\cdot|x,a)$, $\mathcal{T}^{\pi}f$ will not be discrete since here $r(x,a)$
is continuous). Even for regular RL, without Bellman completeness, in the off-
policy setting, it is possible that TD-based algorithms may diverge, and TD
fixed point solutions can be arbitrarily bad.

##  5 Simulation

Figure 1: Visualization of the combination lock. The dotted lines denote
transiting from good states (white) to bad states (gray). Once the agent
transits to a bad state, it stays there forever. The observation is composed
of three parts: one-hot encoding of the latent state $w_{h}$, one-hot encoding
of the time step $h$, and random noise.

In this section, we show the empirical performance of two instances of FLE’s
algorithm: _GMM-FLE_ , which uses Gaussian mixture models (GMM) for
$\mathcal{F}$, and _Diff-FLE_ , which uses diffusion probabilistic models. We
elaborate on each component of the experiments below. See Append D.3 for a
full list of experiment results. We release our code at
https://github.com/ziqian2000/Fitted-Likelihood-Estimation.

##### GMM-FLE.

The GMM-FLE uses conditional Gaussian mixture models, where the weights and
the mean and covariance of Gaussians are all learnable.

##### Diff-FLE.

For Diff-FLE, we simply model the distribution $f(\cdot\,|\,x,a)$ as a
conditional diffusion probabilistic model [Sohl-Dickstein et al., 2015]. Our
implementation is based on DDPM [Ho et al., 2020].

##### The combination lock environment.

There are two chains in the combination lock. Intuitively, the first chain is
a “good” chain, while the second one is bad. The agent wants to stay on the
first chain, for which the only approach is to take the unique optimal action
at each time step. See Figure 1 for an illustration. Mathematically, the
combination lock is a finite horizon MDP of horizon $H$. There are two latent
states $w_{h}\in\\{0,1\\}$. At any time step $h\in[H]$, there is only one
optimal action $a^{\star}_{h}$ among $A$ actions. If the agent is in the
latent state $w_{h}=0$ and takes $a^{\star}_{h}$, it transits to $w_{h+1}=0$,
otherwise transits to $w_{h+1}=1$. If it is already in $w_{h}=1$, no matter
what action it takes, the latent state transits to $w_{h+1}=1$. When $h=H$, it
receives a random reward $r^{+}$ if $w_{H}=0$; otherwise, it gets $r^{-}$. The
agent cannot observe the latent state $w_{h}$ directly. Instead, the
observation it receives, $\psi(w_{h},h)$, is the concatenation of one-hot
coding of the latent state $w_{h}$ and the current time step $h$, appended
with Gaussian noise. This environment has been used in multiple prior works
[Misra et al., 2020, Zhang et al., 2022] where it was shown that standard deep
RL methods struggle due to the challenges from exploration and high-
dimensional observation.

##### Test policy.

The test policy is stochastic: it takes a random action with probability
$\epsilon$ and takes the optimal policy otherwise. In all experiments, we set
$\epsilon=1/7$.

##### Offline data generation.

The offline dataset is generated uniformly. Specifically, for each time step
$h\in[H]$ and each latent state $w_{h}\in\\{0,1\\}$, we first randomly sample
10000 observable state $\phi(w_{h})$. Then for each of them, we uniformly
randomly sample action and perform one step simulation. It is clear that the
offline data distribution here satisfies the coverage assumption (Assumption
4.1).

###  5.1 One-Dimensional Reward

To compare to classic methods such as categorical algorithm [Bellemare et al.,
2017] and quantile TD [Dabney et al., 2018], we first run experiments with a
1-d reward. Specifically, we set $r^{+}\sim\mathcal{N}(1,0.1^{2})$ and
$r^{-}\sim\mathcal{N}(-1,0.1^{2})$. The horizon is $H=20$.

The categorical algorithm discretizes the range $[-1.5,1.5]$ using 100 atoms.
For quantile TD, we set the number of quantiles to 100 as well. The GMM-FLE
uses 10 atomic Gaussian distributions, although eventually, only two are
significant. See Append D for a detailed description of implementations. We
plot the pdfs
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}\hat{f}_{h}(x,a^{\star}_{h})$ (here $0$
denotes the good latent state in $h$) learned by different methods in Figure
2, at three different time steps. As we can see, GMM-FLE in general fits the
ground truth the best.

We also compute the approximated $d_{tv}$ between the learned distribution and
the true one. Ideally, we want to compute
$d_{tv}(\mathop{\mathbb{E}}_{x\sim\psi(0,h)}\hat{f}_{h}(x,a^{\star}_{h}),\mathop{\mathbb{E}}_{x\sim\psi(0,h)}Z^{\pi}_{h}(x,a^{\star}_{h}))$.
However, since obtaining the density of certain models is impossible (e.g.,
Diff-FLE) and certain other models have only discrete supports, we use an
approximated version: we sample $20k$ points from each distribution, construct
two histograms, and calculate $d_{tv}$ between the two histograms. The results
are shown in Table 1. Again, GMM-FLE achieves the smallest total variation
distance. This intuitively makes sense since the ground truth return is a
mixture of Gaussians. Moreover, we notice that GMM-FLE, Diff-FLE, and
categorical algorithms achieve significantly better performance than the
quantile regression TD algorithm. This perhaps is not surprising because our
theory has provided performance guarantees for those three algorithms under
$d_{tv}$ (recall that the categorical algorithm can be roughly considered a
specification of FLE, see Remark 3.1), while it is unclear if quantile
regression TD can achieve similar guarantees in the off-policy learning
setting.

$h$ | Cate Alg | Quan Alg | Diff-FLE | GMM-FLE
---|---|---|---|---
1 |  0.071 $\pm$ 0.015 |  0.603 $\pm$ 0.011 |  0.292 $\pm$ 0.073 | 0.039 $\pm$ 0.004
10 |  0.079 $\pm$ 0.017 |  0.494 $\pm$ 0.018 |  0.234 $\pm$ 0.043 | 0.044 $\pm$ 0.012
19 |  0.078 $\pm$ 0.011 |  0.167 $\pm$ 0.019 |  0.109 $\pm$ 0.031 | 0.018 $\pm$ 0.008
Table 1: Approximated $d_{tv}$ between
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}\hat{f}_{h}(x,a^{\star}_{h})$ and
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}Z^{\pi}_{h}(x,a^{\star}_{h})$ in the 1-d
case. The means and standard errors are computed via five independent runs.

Figure 2: Plots of
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}\hat{f}_{h}(x,a^{\star}_{h})$ and
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}Z^{\pi}_{h}(x,a^{\star}_{h})$. The
histograms are generated via 50k samples.

Figure 3: Plots of
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}\hat{f}_{h}(x,a^{\star}_{h})$ (generated
via 50k samples), and the ground truth
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}Z^{\pi}_{h}(x,a^{\star}_{h})$ (top row).

###  5.2 Two-Dimensional Reward

We conducted experiments on a two-dimensional reward. For now, $r^{+}$ is
sampled from a ring in $\mathbb{R}^{2}$ of radius 2, and $r^{-}$ is a Gaussian
centered at the origin. The horizon is $H=10$. The categorical algorithm
discretizes the range $[-4,4]^{2}$ using $30$ atoms for each dimension (thus
using $900$ atoms in total). Note that the two-dimensional version of the
categorical algorithm is not introduced in the original paper [Bellemare et
al., 2017], but the extension is intuitive. The GMM-FLE uses 30 atomic
Gaussian distributions, although eventually, at most six of them are
significant. We note that it is not straightforward to extend quantile
regression TD to multi-dimensional rewards. More implementation details can be
found in Appendix D.2.

$h$ | Cate Alg | Diff-FLE | GMM-FLE
---|---|---|---
1 |  0.483 $\pm$ 0.003 | 0.357 $\pm$ 0.031 |  0.438 $\pm$ 0.008
5 |  0.466 $\pm$ 0.001 | 0.310 $\pm$ 0.019 |  0.493 $\pm$ 0.050
9 |  0.453 $\pm$ 0.001 | 0.207 $\pm$ 0.014 |  0.502 $\pm$ 0.094
Table 2: Approximated $d_{tv}$ between
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}\hat{f}_{h}(x,a^{\star}_{h})$ and
$\mathop{\mathbb{E}}_{x\sim\psi(0,h)}Z^{\pi}_{h}(x,a^{\star}_{h})$ in the 2-d
case. The means and standard errors are computed via five independent runs.

We plot the 2-d visualization of the learned distribution in Figure 3 and also
computed the approximated TV distance using the same method as in the 1-d
case. The numerical results are shown in Table 2. Note that Diff-FLE performs
the best (Table 2), and as a powerful model, it captures the correlation among
dimensions (i.e., see Figure 3 where Diff-FLE does capture the ring structures
in all steps). However, the GMM-FLE also doesn’t perform well since it is hard
for vanilla GMM with a finite number of mixtures to capture a ring-like data
distribution. The two-dimensional categorical algorithm performed badly as
well, even though it uses a larger number of atoms (recall that for the 1-d
case it only uses 100 atoms and already achieves excellent performance),
implying that it suffers from the curse of dimensionality statistically, i.e.,
explicitly discretizing the 2-d return space evenly can fail to capture the
underlying data structure (e.g., in our ring example, data actually
approximately lives in a submanifold). Moreover, the training is also
significantly slower. In our implementation, we found that running the 2-d
categorical algorithm with $100^{2}$ atoms is about 100 times slower than
running the 1-d algorithm with $100$ atoms, while the training time of Diff-
FLE and GMM-FLE does not change too much.

##  6 Conclusion

In this paper, we proposed a simple algorithm, Fitted Likelihood Estimation
(FLE), for distributional OPE where rewards could be multi-dimensional
vectors. FLE consists of a sequence of MLEs and has the flexibility to
integrate any state-of-art generative models as long as it can be trained via
MLE. Thus, FLE is scalable to the setting where reward vectors are high-
dimensional. Theoretically, we showed that the distribution learned by FLE is
accurate under total variation distance and $p$-Wasserstein distance, for the
finite-horizon setting and the infinite-horizon discounted setting,
respectively. In practice, we show that FLE is flexible enough to use
generative models such as GMMs and diffusion models.

## References

  * Agarwal et al. [2020a] Alekh Agarwal, Mikael Henaff, Sham Kakade, and Wen Sun.  Pc-pg: Policy cover directed exploration for provable policy gradient learning.  _Advances in neural information processing systems_ , 33:13399–13412, 2020a.
  * Agarwal et al. [2020b] 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, 2020b.
  * Bellemare et al. [2017] 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.
  * Bellemare et al. [2023] Marc G. Bellemare, Will Dabney, and Mark Rowland.  _Distributional Reinforcement Learning_.  MIT Press, 2023.  http://www.distributional-rl.org.
  * Chandak et al. [2021] Yash Chandak, Scott Niekum, Bruno da Silva, Erik Learned-Miller, Emma Brunskill, and Philip S Thomas.  Universal off-policy evaluation.  _Advances in Neural Information Processing Systems_ , 34:27475–27490, 2021.
  * Chang et al. [2022] 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.
  * Dabney et al. [2018] 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, 2018.
  * Dinh et al. [2014] Laurent Dinh, David Krueger, and Yoshua Bengio.  Nice: Non-linear independent components estimation.  _arXiv preprint arXiv:1410.8516_ , 2014.
  * Ernst et al. [2005] Damien Ernst, Pierre Geurts, and Louis Wehenkel.  Tree-based batch mode reinforcement learning.  _Journal of Machine Learning Research_ , 6, 2005.
  * Feng et al. [2019] Yihao Feng, Lihong Li, and Qiang Liu.  A kernel loss for solving the bellman equation.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Fu et al. [2021] Justin Fu, Mohammad Norouzi, Ofir Nachum, George Tucker, Ziyu Wang, Alexander Novikov, Mengjiao Yang, Michael R Zhang, Yutian Chen, Aviral Kumar, et al.  Benchmarks for deep off-policy evaluation.  _arXiv preprint arXiv:2103.16596_ , 2021.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Jiang and Li [2016] Nan Jiang and Lihong Li.  Doubly robust off-policy value evaluation for reinforcement learning.  In _International Conference on Machine Learning_ , pages 652–661. PMLR, 2016.
  * Keramati et al. [2020] Ramtin Keramati, Christoph Dann, Alex Tamkin, and Emma Brunskill.  Being optimistic to be conservative: Quickly learning a cvar policy.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pages 4436–4443, 2020.
  * Kolter [2011] J Kolter.  The fixed points of off-policy td.  _Advances in Neural Information Processing Systems_ , 24, 2011.
  * Levin and Peres [2017] David A Levin and Yuval Peres.  _Markov chains and mixing times_ , volume 107.  American Mathematical Soc., 2017.
  * Ma et al. [2021] Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani.  Conservative offline distributional reinforcement learning.  _Advances in Neural Information Processing Systems_ , 34:19235–19247, 2021.
  * Misra et al. [2020] Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford.  Kinematic state abstraction and provably efficient rich-observation reinforcement learning.  In _International conference on machine learning_ , pages 6961–6971. PMLR, 2020.
  * Morimura et al. [2012] Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka.  Parametric return density estimation for reinforcement learning.  _arXiv preprint arXiv:1203.3497_ , 2012.
  * Munos [2003] Rémi Munos.  Error bounds for approximate policy iteration.  In _ICML_ , volume 3, pages 560–567. Citeseer, 2003.
  * Munos and Szepesvári [2008] Rémi Munos and Csaba Szepesvári.  Finite-time bounds for fitted value iteration.  _Journal of Machine Learning Research_ , 9(5), 2008.
  * Precup et al. [2000] Doina Precup, Richard S Sutton, and Satinder P Singh.  Eligibility traces for off-policy policy evaluation.  In _ICML_ , 2000.
  * Rowland et al. [2023] 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 temporal-difference learning.  _arXiv preprint arXiv:2301.04462_ , 2023.
  * Scherrer [2010] Bruno Scherrer.  Should one compute the temporal difference fix point or minimize the bellman residual? the unified oblique projection view.  _arXiv preprint arXiv:1011.4362_ , 2010.
  * Singh and Póczos [2018] Shashank Singh and Barnabás Póczos.  Minimax distribution estimation in wasserstein distance.  _arXiv preprint arXiv:1802.08855_ , 2018.
  * Sohl-Dickstein et al. [2015] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _International Conference on Machine Learning_ , pages 2256–2265. PMLR, 2015.
  * Tsitsiklis and Van Roy [1996] John Tsitsiklis and Benjamin Van Roy.  Analysis of temporal-diffference learning with function approximation.  _Advances in neural information processing systems_ , 9, 1996.
  * Uehara and Sun [2021] Masatoshi Uehara and Wen Sun.  Pessimistic model-based offline reinforcement learning under partial coverage.  _arXiv preprint arXiv:2107.06226_ , 2021.
  * Uehara et al. [2020] Masatoshi Uehara, Jiawei Huang, and Nan Jiang.  Minimax weight and q-function learning for off-policy evaluation.  In _International Conference on Machine Learning_ , pages 9659–9668. PMLR, 2020.
  * Uehara et al. [2021] 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.
  * Van de Geer [2000] Sara Van de Geer.  _Empirical Processes in M-estimation_ , volume 6.  Cambridge university press, 2000.
  * Villani [2021] Cédric Villani.  _Topics in optimal transportation_ , volume 58.  American Mathematical Soc., 2021.
  * Villani et al. [2009] Cédric Villani et al.  _Optimal transport: old and new_ , volume 338.  Springer, 2009.
  * Xie et al. [2021] 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.
  * Yang et al. [2020] Mengjiao Yang, Ofir Nachum, Bo Dai, Lihong Li, and Dale Schuurmans.  Off-policy evaluation via the regularized lagrangian.  _Advances in Neural Information Processing Systems_ , 33:6551–6561, 2020.
  * Zhan et al. [2022] Wenhao Zhan, Masatoshi Uehara, Wen Sun, and Jason D Lee.  Pac reinforcement learning for predictive state representations.  _arXiv preprint arXiv:2207.05738_ , 2022.
  * Zhang et al. [2021] Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, and Tie-Yan Liu.  Distributional reinforcement learning for multi-dimensional reward functions.  _Advances in Neural Information Processing Systems_ , 34:1519–1529, 2021.
  * Zhang [2006] Tong Zhang.  From $\epsilon$-entropy to kl-entropy: Analysis of minimum information complexity density estimation.  _The Annals of Statistics_ , 34(5):2180–2210, 2006.
  * Zhang et al. [2022] Xuezhou Zhang, Yuda Song, Masatoshi Uehara, Mengdi Wang, Alekh Agarwal, and Wen Sun.  Efficient reinforcement learning in block mdps: A model-free representation learning approach.  In _International Conference on Machine Learning_ , pages 26517–26547. PMLR, 2022.

##  Appendix A Missing Proofs in Section 4

###  A.1 Proof of Theorem 4.2

###### Proof.

Note that for all $h\in[H]$, we have

$\displaystyle\mathop{\mathbb{E}}_{x,a\sim d^{\pi}_{h}}\
d_{tv}\left([\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a),[\mathcal{T}^{\pi}Z^{\pi}_{h+1}](x,a)\right)$
$\displaystyle=$ $\displaystyle\frac{1}{2}\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{h}}\sup_{g:\|g\|_{\infty}\leq
1}\left|\mathop{\mathbb{E}}_{\begin{subarray}{c}x^{\prime}\sim P(x,a)\\\
a^{\prime}\sim\pi(x^{\prime})\\\ r\sim
r(x,a)\end{subarray}}\left(\mathop{\mathbb{E}}_{y\sim\hat{f}_{h+1}(\cdot|x^{\prime},a^{\prime})}g\big{(}r+y\big{)}-\mathop{\mathbb{E}}_{y\sim
Z^{\pi}_{h+1}(\cdot|x^{\prime},a^{\prime})}g\big{(}r+y\big{)}\right)\right|$
$\displaystyle\leq$
$\displaystyle\frac{1}{2}\mathop{\mathbb{E}}_{\begin{subarray}{c}x,a\sim
d^{\pi}_{h}\\\ x^{\prime}\sim P(x,a)\\\ a^{\prime}\sim\pi(\cdot|x)\\\ r\sim
r(x,a)\end{subarray}}\sup_{g:\|g\|_{\infty}\leq
1}\left|\mathop{\mathbb{E}}_{y\sim\hat{f}_{h+1}(\cdot|x^{\prime},a^{\prime})}g\big{(}r+y\big{)}-\mathop{\mathbb{E}}_{y\sim
Z^{\pi}_{h+1}(\cdot|x^{\prime},a^{\prime})}g\big{(}r+y\big{)}\right|$
$\displaystyle=$
$\displaystyle\frac{1}{2}\mathop{\mathbb{E}}_{\begin{subarray}{c}x,a\sim
d^{\pi}_{h}\\\ x^{\prime}\sim P(x,a)\\\
a^{\prime}\sim\pi(\cdot|x)\end{subarray}}\sup_{g:\|g\|_{\infty}\leq
1}\left|\mathop{\mathbb{E}}_{y\sim\hat{f}_{h+1}(\cdot|x^{\prime},a^{\prime})}g(y)-\mathop{\mathbb{E}}_{y\sim
Z^{\pi}_{h+1}(\cdot|x^{\prime},a^{\prime})}g(y)\right|$ $\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}x^{\prime},a^{\prime}\sim
d^{\pi}_{h+1}\end{subarray}}d_{tv}\left(\hat{f}_{h+1}(x^{\prime},a^{\prime}),Z^{\pi}_{h+1}(x^{\prime},a^{\prime})\right).$

Here the inequality holds for Jensen’s inequality. The second equality holds
since we can set $\tilde{g}(y)=g(r+y)$ for which we have
$\|\tilde{g}\|_{\infty}\leq 1$ thus removing the additive term $r$. Hence, by
triangle inequality, we have

$\displaystyle\mathop{\mathbb{E}}_{x,a\sim d^{\pi}_{h}}\
d_{tv}\left(\hat{f}_{h}(x,a),Z^{\pi}_{h}(x,a)\right)=\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{h}}\
d_{tv}\left(\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}Z^{\pi}_{h+1}](x,a)\right)$
$\displaystyle\leq$ $\displaystyle\underbrace{\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{h}}\
d_{tv}\left(\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a)\right)}_{\rm(i)}+\underbrace{\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{h}}\
d_{tv}\left([\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a),[\mathcal{T}^{\pi}Z^{\pi}_{h+1}](x,a)\right)\
}_{\rm(ii)}.$

By Assumption 4.1 and Jensen’s inequality, we have
${\rm(i)}\leq\sqrt{C}\zeta_{h}$ because

$\displaystyle\mathop{\mathbb{E}}_{x,a\sim d^{\pi}_{h}}\
d_{tv}\left(\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a)\right)$
$\displaystyle\leq\left\\{\mathop{\mathbb{E}}_{x,a\sim d^{\pi}_{h}}\
d^{2}_{tv}\left(\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a)\right)\right\\}^{1/2}\leq\sqrt{C}\zeta_{h}.$

And by the above derivation we have ${\rm(ii)}\leq\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{h+1}}d_{tv}\left(\hat{f}_{h+1}(x,a),Z^{\pi}_{h+1}(x,a)\right)$.
Hence,

$\mathop{\mathbb{E}}_{x,a\sim d^{\pi}_{h}}\
d_{tv}\left(\hat{f}_{h}(x,a),Z^{\pi}_{h}(x,a)\right)\leq\sqrt{C}\zeta_{h}+\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{h+1}}d_{tv}\left(\hat{f}_{h+1}(x,a),Z^{\pi}_{h+1}(x,a)\right).$

Summing over $h=1,\dots,H$ on both sides, we get

$\mathop{\mathbb{E}}_{x,a\sim d^{\pi}_{1}}\
d_{tv}\left(\hat{f}_{1}(x,a),Z^{\pi}_{1}(x,a)\right)\leq\sqrt{C}\sum_{h=1}^{H}\zeta_{h}+\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{H+1}}d_{tv}\left(\hat{f}_{H+1}(x,a),Z^{\pi}_{H+1}(x,a)\right)=\sqrt{C}\sum_{h=1}^{H}\zeta_{h}.$

where the equality holds since $\hat{f}_{H+1}=Z^{\pi}_{H+1}=0$ by definition.
Now we complete the proof by noticing the following

$\displaystyle
d_{tv}\left(\hat{f},Z^{\pi}\right)=\frac{1}{2}\sup_{g:\|g\|_{\infty}\leq
1}\left|\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{1}}\left(\mathop{\mathbb{E}}_{y\sim\hat{f}_{1}(\cdot|x,a)}g(y)-\mathop{\mathbb{E}}_{y\sim
Z^{\pi}(\cdot|x,a)}g(y)\right)\right|$ $\displaystyle\leq$
$\displaystyle\frac{1}{2}\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}_{1}}\sup_{g:\|g\|_{\infty}\leq
1}\left|\mathop{\mathbb{E}}_{y\sim\hat{f}_{1}(\cdot|x,a)}g(y)-\mathop{\mathbb{E}}_{y\sim
Z^{\pi}(\cdot|x,a)}g(y)\right|=\mathop{\mathbb{E}}_{x,a\sim d^{\pi}_{1}}\
d_{tv}\left(\hat{f}_{1}(x,a),Z^{\pi}_{1}(x,a)\right).$

∎

###  A.2 Proof of Lemma 4.4

###### Proof.

Observing Algorithm 1, when $h=H$, we are estimating the conditional
distribution $Z^{\pi}_{H}$ via MLE. Under Assumption 4.3 which implies that
there exists a function $g\in\mathcal{F}_{H}$ that agrees with $Z^{\pi}_{H}$
on the support of $\rho$, we can apply Lemma C.3, which leads to

$\mathop{\mathbb{E}}_{x,a\sim\rho}d_{tv}^{2}\left(\hat{f}_{H}(x,a),Z^{\pi}_{H}(x,a)\right)\leq\frac{4H}{n}\log(|\mathcal{F}_{H}|/\delta)$

with probability at least $1-\delta$. When $h<H$, we are estimating the
conditional distribution $\mathcal{T}^{\pi}\hat{f}_{h+1}$ via MLE. Also note
that thanks to the random data split, we have $\hat{f}_{h+1}$ being
independent of the dataset $\mathcal{D}_{h}$ ($\hat{f}_{h+1}$ only depends on
datasets $\mathcal{D}_{h+1},\dots\mathcal{D}_{H}$). Therefore, under
Assumption 4.3 which implies that there exists a function
$g\in\mathcal{F}_{h}$ that agrees with $\mathcal{T}^{\pi}\hat{f}_{h+1}$ on the
support of $\rho$, we can apply Lemma C.3, which leads to

$\mathop{\mathbb{E}}_{x,a\sim\rho}d_{tv}^{2}\left(\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a)\right)\leq\frac{4H}{n}\log(|\mathcal{F}_{H}|/\delta)$

with probability at least $1-\delta$. We complete the proof by taking the
union bound for $h\in[H]$.

∎

###  A.3 Proof of Lemma 4.6

###### Proof.

The proof is similar to Lemma 4.4. Observing Algorithm 1, when $h=H$, we are
basically estimating the conditional distribution $Z^{\pi}_{H}$ via MLE.
Hence, under Assumption 4.3 which implies that there exists a function
$g\in\mathcal{F}_{H}$ that agrees with $Z^{\pi}_{H}$ on the support of $\rho$,
we can apply Lemma C.3, which leads to

$\mathop{\mathbb{E}}_{x,a\sim\rho}d_{tv}^{2}\left(\hat{f}_{H}(x,a),Z^{\pi}_{H}(x,a)\right)\leq\frac{10H}{n}\log\left(N_{[]}\big{(}(nH^{d})^{-1},\mathcal{F}_{H},\|\cdot\|_{\infty}\big{)}/\delta\right)$

with probability at least $1-\delta$. When $h<H$, we are estimating the
conditional distribution $\mathcal{T}^{\pi}\hat{f}_{h+1}$ via MLE. Therefore,
under Assumption 4.3 which implies that there exists a function
$g\in\mathcal{F}_{h}$ that agrees with $\mathcal{T}^{\pi}\hat{f}_{h+1}$ on the
support of $\rho$, we can apply Lemma C.3, which leads to

$\mathop{\mathbb{E}}_{x,a\sim\rho}d_{tv}^{2}\left(\hat{f}_{h}(x,a),[\mathcal{T}^{\pi}\hat{f}_{h+1}](x,a)\right)\leq\frac{10H}{n}\log\left(N_{[]}\big{(}(nH^{d})^{-1},\mathcal{F}_{h},\|\cdot\|_{\infty}\big{)}/\delta\right)$

with probability at least $1-\delta$. We complete the proof by taking the
union bound for $h\in[H]$. ∎

###  A.4 Proof of Lemma 4.9

###### Proof.

First, it deserves to verify that the “metric” $(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p})^{1/(2p)}$ we are using satisfies the triangle inequality
and is thus indeed a metric. To this end, we note that, for any three
densities
$f_{1},f_{2},f_{3}:\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,(1-\gamma)^{-1}]^{d})$,
the folloing holds since $d_{w,p}$ is a metric,

$\displaystyle\bigg{(}\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}(f_{1}(x,a),f_{2}(x,a))\bigg{)}^{\frac{1}{2p}}\leq\bigg{(}\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}\Big{(}d_{w,p}(f_{1}(x,a),f_{3}(x,a))+d_{w,p}(f_{3}(x,a),f_{2}(x,a))\Big{)}^{2p}\bigg{)}^{\frac{1}{2p}}.$

Then by Minkowski inequality, the above

$\displaystyle\leq\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}(f_{1}(x,a),f_{3}(x,a))\right)^{\frac{1}{2p}}+\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}(f_{3}(x,a),f_{2}(x,a))\right)^{\frac{1}{2p}},$

for which we conclude triangle inequality for $(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p})^{1/(2p)}$. Since other axioms of metrics are trivial to
verify, we conclude that it is indeed a metric. Hence, we can safely proceed.

To establish the contractive property, we start with the following lemma,
which shows that the distributional Bellman operator is roughly
“$\gamma$-contractive” in a sense but with distribution shifts.

######  Lemma A.1.

For any $f,f^{\prime}\in\mathcal{F}$, $x\in\mathcal{X}$ and $a\in\mathcal{A}$,
we have

$\displaystyle
d_{w,p}^{p}\left([\mathcal{T}^{\pi}f](x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)\leq\mathop{\mathbb{E}}_{x^{\prime}\sim
P(x,a),a^{\prime}\sim\pi(x^{\prime})}\gamma^{p}d^{p}_{w,p}\left(f(x^{\prime},a^{\prime}),f^{\prime}(x^{\prime},a^{\prime})\right).$

###### Proof of Lemma A.1.

By the dual form of Wasserstein distance (Lemma C.5), we have

$\displaystyle
d_{w,p}^{p}\left([\mathcal{T}^{\pi}f](x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)=\sup_{(\psi,\phi)\in\Gamma}\mathop{\mathbb{E}}_{z\sim[\mathcal{T}^{\pi}f](x,a)}\psi(z)-\mathop{\mathbb{E}}_{z\sim[\mathcal{T}^{\pi}f^{\prime}](x,a)}\phi(z)$
$\displaystyle=$
$\displaystyle\sup_{(\psi,\phi)\in\Gamma}\mathop{\mathbb{E}}_{\begin{subarray}{c}x^{\prime}\sim
P(x,a)\\\ a^{\prime}\sim\pi(x^{\prime})\\\ r\sim
r(x,a)\end{subarray}}\left(\mathop{\mathbb{E}}_{y\sim
f(x^{\prime},a^{\prime})}\psi(r+\gamma y)-\mathop{\mathbb{E}}_{y\sim
f^{\prime}(x^{\prime},a^{\prime})}\phi(r+\gamma y)\right)$ $\displaystyle\leq$
$\displaystyle\mathop{\mathbb{E}}_{\begin{subarray}{c}x^{\prime}\sim P(x,a)\\\
a^{\prime}\sim\pi(x^{\prime})\\\ r\sim
r(x,a)\end{subarray}}\underbrace{\sup_{(\psi,\phi)\in\Gamma}\left(\mathop{\mathbb{E}}_{y\sim
f(x^{\prime},a^{\prime})}\psi(r+\gamma y)-\mathop{\mathbb{E}}_{y\sim
f^{\prime}(x^{\prime},a^{\prime})}\phi(r+\gamma y)\right)}_{(*)}$ (3)

where $\Gamma=\\{(\psi,\phi):\psi(x)-\phi(y)\leq\|x-y\|^{p}\\}$. The second
equality holds by the definition of Bellman operator.

Regarding $\rm(*)$, for any $(\psi,\phi)\in\Gamma$, we define
$\widetilde{\psi}(y)=\psi(r+\gamma y)/\gamma^{p}$ and
$\widetilde{\phi}(y)=\phi(r+\gamma y)/\gamma^{p}$. Then, we have

$\displaystyle(*)=\gamma^{p}\sup_{(\psi,\phi)\in\Gamma}\left(\mathop{\mathbb{E}}_{y\sim
f(x^{\prime},a^{\prime})}\widetilde{\psi}(y)-\mathop{\mathbb{E}}_{y\sim
f^{\prime}(x^{\prime},a^{\prime})}\widetilde{\phi}(y)\right).$

We note that, for any $x,y$,

$\displaystyle\widetilde{\psi}(x)-\widetilde{\phi}(y)=\frac{\psi(r+\gamma
x)-\phi(r+\gamma y)}{\gamma^{p}}\leq\frac{\|(r+\gamma x)-(r+\gamma
y)\|^{p}}{\gamma^{p}}=\|x-y\|^{p}.$

Here the inequality holds since $(\psi,\phi)\in\Gamma$. Hence,
$(\widetilde{\psi},\widetilde{\phi})\in\Gamma$ as well. In other words, for
any given $\psi$ and $\phi$, their correspondences $\widetilde{\psi}$ and
$\widetilde{\phi}$ are also in $\Gamma$. Thus we can take the supremum
directly over the latter, which leads to

$\displaystyle(*)\leq\gamma^{p}\sup_{(\widetilde{\psi},\widetilde{\phi})\in\Gamma}\left(\mathop{\mathbb{E}}_{y\sim
f(x^{\prime},a^{\prime})}\widetilde{\psi}(y)-\mathop{\mathbb{E}}_{y\sim
f^{\prime}(x^{\prime},a^{\prime})}\widetilde{\phi}(y)\right)=\gamma^{p}d^{p}_{w,p}\big{(}f(x^{\prime},a^{\prime}),f^{\prime}(x^{\prime},a^{\prime})\big{)}$

where the equality holds due to the dual form of Wasserstein distance (Lemma
C.5) again. Then we plug the above into (3) and get

$\displaystyle
d_{w,p}^{p}\left([\mathcal{T}^{\pi}f](x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)\leq\mathop{\mathbb{E}}_{\begin{subarray}{c}x^{\prime}\sim
P(x,a),a^{\prime}\sim\pi(x^{\prime})\end{subarray}}\gamma^{p}d^{p}_{w,p}\big{(}f(x^{\prime},a^{\prime}),f^{\prime}(x^{\prime},a^{\prime})\big{)}.$

where we have removed the randomness of $r\sim r(x,a)$ originally appeared in
(3) since the the term inside the expectation is now completely independent of
$r$. ∎

By Lemma A.1, we have

$\displaystyle\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left([\mathcal{T}^{\pi}f](x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)\right)^{\frac{1}{2p}}=\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}\Big{(}d^{p}_{w,p}\left([\mathcal{T}^{\pi}f](x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)\Big{)}^{2}\right)^{\frac{1}{2p}}$
$\displaystyle\leq$
$\displaystyle\gamma\cdot\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}\left(\mathop{\mathbb{E}}_{x^{\prime}\sim
P(x,a),a^{\prime}\sim\pi(x^{\prime})}d^{p}_{w,p}\left(f(x^{\prime},a^{\prime}),f^{\prime}(x^{\prime},a^{\prime})\right)\right)^{2}\right)^{\frac{1}{2p}}$
$\displaystyle\leq$
$\displaystyle\gamma\cdot\left(\underbrace{\mathop{\mathbb{E}}_{\begin{subarray}{c}x,a\sim
d^{\pi}\\\ x^{\prime}\sim
P(x,a),a^{\prime}\sim\pi(x^{\prime})\end{subarray}}d^{2p}_{w,p}\left(f(x^{\prime},a^{\prime}),f^{\prime}(x^{\prime},a^{\prime})\right)}_{(\dagger)}\right)^{\frac{1}{2p}}$

where the last inequality holds because of Jensen’s inequality. Since
$d^{\pi}(x,a)=\gamma\mathop{\mathbb{E}}_{\tilde{x},\tilde{a}\sim
d^{\pi}}P(x|\tilde{x},\tilde{a})\pi(a|x)+(1-\gamma)\mu(x)\pi(x|a)$, we have
$\mathop{\mathbb{E}}_{\tilde{x},\tilde{a}\sim
d^{\pi}}P(x|\tilde{x},\tilde{a})\pi(a|x)\leq\gamma^{-1}d^{\pi}(x,a)$.
Therefore,

$\displaystyle(\dagger)\leq\gamma^{-1}\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(f(x,a),f^{\prime}(x,a)\right).$

Hence, we conclude that

$\displaystyle\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left([\mathcal{T}^{\pi}f](x,a),[\mathcal{T}^{\pi}f^{\prime}](x,a)\right)\right)^{\frac{1}{2p}}$
$\displaystyle\leq$
$\displaystyle\gamma\cdot\left(\gamma^{-1}\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(f(x,a),f^{\prime}(x,a)\right)\right)^{\frac{1}{2p}}$
$\displaystyle=$
$\displaystyle\gamma^{1-\frac{1}{2p}}\cdot\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(f(x,a),f^{\prime}(x,a)\right)\right)^{\frac{1}{2p}}.$

∎

###  A.5 Proof of Theorem 4.11

###### Proof.

We will prove the following theorem which is more general.

######  Theorem A.2.

Under 4.10, suppose we have a sequence of functions
$\hat{f}_{1},\dots,\hat{f}_{T}:\mathcal{X}\times\mathcal{A}\mapsto\Delta([0,(1-\gamma)^{-1}]^{d})$
and a sequence of values $\zeta_{1},\dots,\zeta_{T}\in\mathbb{R}$ such that

$\textstyle\bigg{(}\mathop{\mathbb{E}}_{x,a\sim\rho}\
d_{w,p}^{2p}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right)\bigg{)}^{\frac{1}{2p}}\leq\zeta_{t}$

holds for all $t\in[T]$. Let our estimator
$\hat{f}\coloneqq\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\hat{f}_{T}(x,a)$.
Then we have, for all $p\geq 1$,

$d_{w,p}\left(\hat{f},Z^{\pi}\right)\leq\left(\frac{C}{1-\gamma}\right)^{\frac{1}{2p}}\sum_{t=1}^{T}\gamma^{(T-t)\left(1-\frac{1}{2p}\right)}\cdot\zeta_{t}+\frac{\sqrt{d}\cdot\gamma^{T\left(1-\frac{1}{2p}\right)}}{(1-\gamma)^{1+\frac{1}{2p}}}.$
(4)

###### Proof of Theorem A.2.

Recall that we defined the conditional distribuions
$\bar{Z}^{\pi}(x,a)\in\Delta([0,(1-\gamma)^{-1}]^{d})$ which is the
distribution of the return under policy $\pi$ starting with state action
$(x,a)$. It is easy to see that
$Z^{\pi}=\mathbb{E}_{x\sim\mu,a\sim\pi(x)}\left[\bar{Z}^{\pi}(x,a)\right]$. We
start with the following.

$\displaystyle\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(\hat{f}_{t}(x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2p}}$
$\displaystyle\leq$ $\displaystyle\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right)\right)^{\frac{1}{2p}}+\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left([\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2p}}$
$\displaystyle\leq$ $\displaystyle\
C^{\frac{1}{2p}}\left(\mathop{\mathbb{E}}_{x,a\sim\rho}d^{2p}_{w,p}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right)\right)^{\frac{1}{2p}}+\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left([\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a),[\mathcal{T}^{\pi}\bar{Z}^{\pi}](x,a)\right)\right)^{\frac{1}{2p}}$
$\displaystyle\leq$ $\displaystyle
C^{\frac{1}{2p}}\zeta_{t}+\gamma^{1-\frac{1}{2p}}\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(\hat{f}_{t-1}(x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2p}}$

where the first inequality is due to triangle inequality (proved in Section
A.4), the second inequality holds because of the coverage assumption (4.10),
and the last inequality holds due to the contractive property of the
distributional Bellman operator (Lemma 4.9). Unrolling the recursion of $t$,
we arrive at

$\displaystyle\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(\hat{f}_{T}(x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2p}}$
$\displaystyle\leq$
$\displaystyle\sum_{t=1}^{T}\gamma^{(T-t)\left(1-\frac{1}{2p}\right)}C^{\frac{1}{2p}}\zeta_{t}+\gamma^{T\left(1-\frac{1}{2p}\right)}\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(\hat{f}_{0}(x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2p}}$
$\displaystyle\leq$
$\displaystyle\sum_{t=1}^{T}\gamma^{(T-t)\left(1-\frac{1}{2p}\right)}C^{\frac{1}{2p}}\zeta_{t}+\gamma^{T\left(1-\frac{1}{2p}\right)}\cdot\frac{\sqrt{d}}{1-\gamma}$
(5)

where the last inequality is due to Corollary C.8 which shows that

$d_{w,p}(\hat{f}_{0}(x,a),\bar{Z}^{\pi}(x,a))\leq\text{\rm
diam}\big{(}[0,(1-\gamma)^{-1}]^{d}\big{)}\leq\frac{\sqrt{d}}{(1-\gamma)}.$

Since $d^{\pi}(x,a)=\gamma\mathop{\mathbb{E}}_{\tilde{x},\tilde{a}\sim
d^{\pi}}P(x|\tilde{x},\tilde{a})\pi(a|x)+(1-\gamma)\mu(x)\pi(x\,|\,a)$, we
have $\mu(x)\pi(x\,|\,a)\leq(1-\gamma)^{-1}d^{\pi}(x,a)$ and thus

$\displaystyle\left(\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}d^{2p}_{w,p}\left(\hat{f}_{T}(x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2p}}\leq\left((1-\gamma)^{-1}\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(\hat{f}_{T}(x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2p}}$
$\displaystyle=$
$\displaystyle(1-\gamma)^{-\frac{1}{2p}}\left(\mathop{\mathbb{E}}_{x,a\sim
d^{\pi}}d^{2p}_{w,p}\left(\hat{f}_{T}(x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2p}}\leq(1-\gamma)^{-\frac{1}{2p}}\left(\sum_{t=1}^{T}\gamma^{(T-t)\left(1-\frac{1}{2p}\right)}C^{\frac{1}{2p}}\zeta_{t}+\gamma^{T\left(1-\frac{1}{2p}\right)}\cdot\frac{\sqrt{d}}{1-\gamma}\right)$
(6)

where the last inequality is for (5).

Applying the dual representation of Wasserstein distance (Lemma C.5) to
$d_{w,p}^{p}\left(\hat{f},Z^{\pi}\right)$, we have

$\displaystyle
d_{w,p}^{p}\left(\hat{f},Z^{\pi}\right)=d_{w,p}^{p}\left(\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\hat{f}_{T}(x,a),\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\bar{Z}^{\pi}(x,a)\right)$
$\displaystyle=$
$\displaystyle\sup_{\psi,\phi\in\Gamma}\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\left(\mathop{\mathbb{E}}_{z\sim\hat{f}_{T}(x,a)}\psi(z)-\mathop{\mathbb{E}}_{z\sim\bar{Z}^{\pi}(x,a)}\phi(z)\right)$
$\displaystyle\leq$
$\displaystyle\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}\sup_{\psi,\phi\in\Gamma}\left(\mathop{\mathbb{E}}_{z\sim\hat{f}_{T}(x,a)}\psi(z)-\mathop{\mathbb{E}}_{z\sim\bar{Z}^{\pi}(x,a)}\phi(z)\right)$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}d^{p}_{w,p}\left(\hat{f}_{T}(x,a),\bar{Z}^{\pi}(x,a)\right)$
$\displaystyle\leq$
$\displaystyle\left(\mathop{\mathbb{E}}_{x\sim\mu,a\sim\pi(x)}d^{2p}_{w,p}\left(\hat{f}_{T}(x,a),\bar{Z}^{\pi}(x,a)\right)\right)^{\frac{1}{2}}.$
(7)

where $\Gamma=\\{(\psi,\phi):\psi(x)-\phi(y)\leq\|x-y\|^{p}\\}$. By chaining
(6) and (7) we complete the proof. ∎

By assuming there exists a common upper bound $\zeta$ (i.e.,
$\zeta_{t}\leq\zeta$, $\forall t$), we can further simplify (4) by noticing
the following. First, since the sum of geometric series is bounded in the
following sense

$\sum_{t=1}^{T}\gamma^{(T-t)\left(1-\frac{1}{2p}\right)}\leq\frac{1}{1-\gamma^{\left(1-\frac{1}{2p}\right)}},$

we can get

$d_{w,p}\left(\hat{f},Z^{\pi}\right)\leq\left(\frac{C}{1-\gamma}\right)^{\frac{1}{2p}}\frac{\zeta}{\left(1-\gamma^{1-\frac{1}{2p}}\right)}+\frac{\sqrt{d}\cdot\gamma^{T\left(1-\frac{1}{2p}\right)}}{(1-\gamma)^{1+\frac{1}{2p}}}.$

Second, we note that the right-hand side above attains the maximum when $p=1$.
Therefore

$\displaystyle d_{w,p}\left(\hat{f},Z^{\pi}\right)\leq$
$\displaystyle\left(\frac{C}{1-\gamma}\right)^{\frac{1}{2p}}\cdot\frac{\zeta}{1-\gamma^{\frac{1}{2}}}+\frac{\sqrt{d}\cdot\gamma^{\frac{T}{2}}}{(1-\gamma)^{\frac{3}{2}}}$
$\displaystyle\leq$
$\displaystyle\frac{2C^{\frac{1}{2p}}}{(1-\gamma)^{\frac{3}{2}}}\cdot\zeta+\frac{\sqrt{d}\cdot\gamma^{\frac{T}{2}}}{(1-\gamma)^{\frac{3}{2}}}$

where the last inequality holds since $1-\gamma^{1/2}\geq(1-\gamma)/2$. ∎

###  A.6 Proof of Lemma 4.13

###### Proof.

We only show the proof for finite function class since the proof for infinite
class is essentially the same.

For Algorithm 2, we are iteratively estimating the conditional distribution
$\mathcal{T}^{\pi}\hat{f}_{t-1}$. Note that thanks to the random data split,
we have $\hat{f}_{t-1}$ being independent of the dataset $\mathcal{D}_{t}$
($\hat{f}_{t-1}$ only depends on datasets
$\mathcal{D}_{1},\dots\mathcal{D}_{t-1}$). Therefore, under Assumption 4.12
which implies that there exists a function $g\in\mathcal{F}$ that agrees with
$\mathcal{T}^{\pi}\hat{f}_{t-1}$ on the support of $\rho$, we can apply Lemma
C.3, which leads to

$\mathop{\mathbb{E}}_{x,a\sim\rho}d_{tv}^{2}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right)\leq\frac{4T}{n}\log(|\mathcal{F}|T/\delta)$

with probability at least $1-\delta$. Here we have taken the union bound for
$t\in[T]$. For the result of Wasserstein distance, we apply Corollary C.7 and
get

$\displaystyle\mathop{\mathbb{E}}_{x,a\sim\rho}d_{w,p}^{2p}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right)\leq\left(\frac{\sqrt{d}}{1-\gamma}\right)^{2p}\mathop{\mathbb{E}}_{x,a\sim\rho}d_{tv}^{2}\left(\hat{f}_{t}(x,a),[\mathcal{T}^{\pi}\hat{f}_{t-1}](x,a)\right).$

∎

###  A.7 Proof of Corollary 4.14

###### Proof.

We only prove for the finite function class ($|\mathcal{F}|<\infty$) since the
proof for the infinite function class is quite similar.

We start with Theorem 4.11, plug in the result of Lemma 4.13, and get

$\displaystyle d_{w,p}\left(\hat{f},Z^{\pi}\right)\leq$
$\displaystyle\frac{2C^{\frac{1}{2p}}}{(1-\gamma)^{\frac{3}{2}}}\cdot\frac{\sqrt{d}}{1-\gamma}\cdot\left(\frac{4T}{n}\log(|\mathcal{F}|T/\delta)\right)^{\frac{1}{2p}}+\frac{\sqrt{d}\cdot\gamma^{\frac{T}{2}}}{(1-\gamma)^{\frac{3}{2}}}$
$\displaystyle=$
$\displaystyle\frac{\sqrt{d}}{(1-\gamma)^{\frac{3}{2}}}\left(\frac{2C^{\frac{1}{2p}}}{1-\gamma}\cdot\left(\frac{4T}{n}\log(|\mathcal{F}|T/\delta)\right)^{\frac{1}{2p}}+\gamma^{\frac{T}{2}}\right)$
(8)

We choose

$\displaystyle
T=\frac{\log\left(C^{\frac{1}{2p}}\cdot\iota^{\frac{1}{2p}}\cdot\left(1-\gamma\right)^{-1}\cdot
n^{-\frac{1}{2p}}\right)}{\log\left(\gamma^{\frac{1}{2}}\right)}\text{\quad
where\quad}\iota=\log(|\mathcal{F}|/\delta),$

which leads to

$\displaystyle\gamma^{\frac{T}{2}}=\frac{C^{\frac{1}{2p}}\cdot\iota^{\frac{1}{2p}}\cdot
n^{-\frac{1}{2p}}}{1-\gamma}.$

Thus, the second additive term of (8) will be smaller than the first one.
Hence, we conclude that

$d_{w,p}\left(\hat{f},Z^{\pi}\right)\leq
2\cdot\frac{\sqrt{d}}{(1-\gamma)^{\frac{3}{2}}}\cdot\frac{2C^{\frac{1}{2p}}}{1-\gamma}\cdot\left(\frac{4T}{n}\log(|\mathcal{F}|T/\delta)\right)^{\frac{1}{2p}}\leq\widetilde{O}\left(\frac{\sqrt{d}\big{(}C\log(|\mathcal{F}|T/\delta)\big{)}^{\frac{1}{2p}}}{(1-\gamma)^{\frac{5}{2}}\cdot
n^{\frac{1}{2p}}}\right).$

∎

###  A.8 Proof of Lemma 4.15

###### Proof.

The bracketing number of the probability simplex
$\Delta(|\mathcal{X}||\mathcal{A}|)$ is bounded by
$N_{[]}(\epsilon,\Delta(|\mathcal{X}||\mathcal{A}|),\|\cdot\|_{\infty})\leq(c/\epsilon)^{|\mathcal{X}||\mathcal{A}|}$
where $c$ is a constant. Hence, we have
$N_{[]}(\epsilon,(\Delta(|\mathcal{X}||\mathcal{A}|))^{|\mathcal{X}||\mathcal{A}|},\|\cdot\|_{\infty})\leq(c/\epsilon)^{|\mathcal{X}|^{2}|\mathcal{A}|^{2}}$.

Let $\widetilde{\Delta}$ denote an $\epsilon$-bracket of
$(\Delta(|\mathcal{X}||\mathcal{A}|))^{|\mathcal{X}||\mathcal{A}|}$. Then we
can construct a bracket of $\mathcal{F}_{h}$ as follows

$\widetilde{\mathcal{F}}_{h}=\left\\{\Big{[}\underline{f},\overline{f}\Big{]}:\underline{f}(x,a)=\sum_{x^{\prime},a^{\prime}}\underline{w}_{x,a}(x^{\prime},a^{\prime})r_{H}(x^{\prime},a^{\prime}),\
\overline{f}(x,a)=\sum_{x^{\prime},a^{\prime}}\overline{w}_{x,a}(x^{\prime},a^{\prime})r_{H}(x^{\prime},a^{\prime}),\
\forall[\underline{w},\overline{w}]\in\widetilde{\Delta}\right\\}.$

We claim that $\widetilde{\mathcal{F}}_{h}$ is a $\epsilon
r_{\infty}|\mathcal{X}||\mathcal{A}|$-bracket of $\mathcal{F}_{h}$. To see
this, we have

$\displaystyle\Big{\|}\overline{f}-\underline{f}\Big{\|}_{\infty}\leq\sum_{x^{\prime},a^{\prime}}\left|\underline{w}_{x,a}(x^{\prime},a^{\prime})-\overline{w}_{x,a}(x^{\prime},a^{\prime})\right|r_{H}(x^{\prime},a^{\prime})\leq\epsilon\sum_{x^{\prime},a^{\prime}}r_{H}(x^{\prime},a^{\prime})\leq\epsilon
r_{\infty}|\mathcal{X}||\mathcal{A}|.$

Therefore, we conclude $N_{[]}(\epsilon
r_{\infty}|\mathcal{X}||\mathcal{A}|,\mathcal{F}_{h},\|\cdot\|_{\infty})\leq|\widetilde{\Delta}|\leq(c/\epsilon)^{|\mathcal{X}|^{2}|\mathcal{A}|^{2}}$.
By substitution we arrive at
$N_{[]}(\epsilon,\mathcal{F}_{h},\|\cdot\|_{\infty})\leq(cr_{\infty}|\mathcal{X}||\mathcal{A}|/\epsilon)^{|\mathcal{X}|^{2}|\mathcal{A}|^{2}}$.
Then we complete the proof by taking a logatithm. ∎

###  A.9 Proof of Lemma 4.16

$\displaystyle\mu_{h}(x,a)=$
$\displaystyle\sum_{i=h}^{H}-(x_{i}^{\top}Qx_{i}+a_{i}^{\top}Ra_{i})=-x^{\top}Qx-a^{\top}Ra-\sum_{i=h+1}^{H}-(x_{i}^{\top}Qx_{i}+a_{i}^{\top}Ra_{i})$
$\displaystyle=$
$\displaystyle-x^{\top}Qx-a^{\top}Ra-\sum_{i=h+1}^{H}(x_{i}^{\top}Qx_{i}+x_{i}^{\top}K^{\top}RKx_{i})$
$\displaystyle=$
$\displaystyle-x^{\top}Qx-a^{\top}Ra-\sum_{i=h+1}^{H}x_{i}^{\top}\left(Q+K^{\top}RK\right)x_{i}$
$\displaystyle=$
$\displaystyle-x^{\top}Qx-a^{\top}Ra-\sum_{i=h+1}^{H}\left((A+BK)^{i-h-1}(Ax+Ba)\right)^{\top}\left(Q+K^{\top}RK\right)\left((A+BK)^{i-h-1}(Ax+Ba)\right)$
$\displaystyle=$
$\displaystyle-x^{\top}Qx-a^{\top}Ra-(Ax+Ba)^{\top}\left(\sum_{i=h+1}^{H}\left((A+BK)^{i-h-1}\right)^{\top}\left(Q+K^{\top}RK\right)(A+BK)^{i-h-1}\right)(Ax+Ba).$

###  A.10 Proof of Lemma 4.17

######  Lemma A.3.

For any $x,a,b\in\mathbb{R}$, we have
$\exp(-(x-a)^{2})-\exp(-(x-b)^{2})\leq\sqrt{2/e}\cdot|a-b|$.

###### Proof of Lemma A.3.

When $a\geq b$, it is equivalent to
$\exp(-(x-a)^{2})-\exp(-(x-b)^{2})\leq\sqrt{2/e}\cdot(a-b)$. Thus it suffices
to show that $g(x,a)\coloneqq\exp(-(x-a)^{2})-\sqrt{2/e}\cdot a$ is non-
increasing in $a$. We take the first derivative with respect to $a$ and then
get

$\frac{\partial}{\partial
a}g(x,a)=2(x-a)\exp\big{(}-(x-a)^{2}\big{)}-\sqrt{\frac{2}{e}}\leq 0$

since it is easy to verify that $\max_{x}|x\exp(-x^{2})|\leq 1/\sqrt{2e}$.
This completes the proof for $a\geq b$.

When $a<b$, it suffices to show that
$h(x,a)\coloneqq\exp(-(x-a)^{2})+\sqrt{2/e}\cdot a$ is non-decreasing in $a$.
We take the first derivative with respect to $a$ and then get

$\frac{\partial}{\partial
a}h(x,a)=2(x-a)\exp\big{(}-(x-a)^{2}\big{)}+\sqrt{\frac{2}{e}}\geq 0.$

Thus we are done. ∎

######  Lemma A.4.

For any $\mu_{1},\mu_{2}\in\mathbb{R}$, it holds that
$\max_{x}\mathcal{N}(x\,|\,\mu_{1},\sigma^{2})-\mathcal{N}(x\,|\,\mu_{2},\sigma^{2})\leq\frac{1}{\sigma^{2}\sqrt{2\pi
e}}\cdot|\mu_{1}-\mu_{2}|$.

###### Proof of Lemma A.4.

$\displaystyle\mathcal{N}(x\,|\,\mu_{1},\sigma^{2})-\mathcal{N}(x\,|\,\mu_{2},\sigma^{2})=\frac{1}{\sigma\sqrt{2\pi}}\left(\exp\left(-\frac{1}{2}\left(\frac{x-\mu_{1}}{\sigma}\right)^{2}\right)-\exp\left(-\frac{1}{2}\left(\frac{x-\mu_{2}}{\sigma}\right)^{2}\right)\right)$
$\displaystyle\leq$
$\displaystyle\frac{1}{\sigma\sqrt{2\pi}}\cdot\sqrt{\frac{2}{e}}\cdot\left|\frac{\mu_{1}}{\sigma\sqrt{2}}-\frac{\mu_{2}}{\sigma\sqrt{2}}\right|=\frac{1}{\sigma^{2}\sqrt{2\pi
e}}|\mu_{1}-\mu_{2}|$

where the inequality holds for Lemma A.3. ∎

######  Lemma A.5.

For LQR, let $\mathcal{M}_{i}$ ($i=1,2,3$) denotes the set of possible
matrices of $M_{i}$. We assume that, there exists parameters $m_{x}$ and
$m_{a}$ for which $\|x\|_{2}\leq m_{x}$ and $\|a\|_{2}\leq m_{a}$ for all
$x\in\mathcal{X}$ and $a\in\mathcal{A}$. Then we have

$\displaystyle N_{[]}$
$\displaystyle(\epsilon,\mathcal{F}_{h},\|\cdot\|_{\infty})\leq\prod_{i=1,2,3}N\left(\frac{\epsilon\sigma^{2}(H-h+1)\sqrt{2\pi
e}}{2(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})},\mathcal{M}_{i},\|\cdot\|_{\rm
F}\right).$

Here $N_{[]}()$ and $N()$ denote the bracketing number and covering number,
respectively.

###### Proof of Lemma A.5.

We denote by $\widetilde{\mathcal{M}}_{1}$, $\widetilde{\mathcal{M}}_{2}$, and
$\widetilde{\mathcal{M}}_{3}$ the $\epsilon$-covers of $\mathcal{M}_{1}$,
$\mathcal{M}_{2}$, and $\mathcal{M}_{3}$, respectively. We construct the
following function class

$\displaystyle\widetilde{\mathcal{F}}_{h}=\Big{\\{}\tilde{f}(\cdot|x,a)=\mathcal{N}\big{(}\cdot\,\big{|}\,x^{\top}\widetilde{M}_{1}x+a^{\top}\widetilde{M}_{2}x+a^{\top}\widetilde{M}_{3}a,(H-h+1)\sigma^{2}\big{)},\
\forall\widetilde{M}_{1}\in\widetilde{\mathcal{M}}_{1},\widetilde{M}_{2}\in\widetilde{\mathcal{M}}_{2},\widetilde{M}_{3}\in\widetilde{\mathcal{M}}_{3}\Big{\\}}.$

We claim that $\widetilde{\mathcal{F}}_{h}$ is a cover of $\mathcal{F}_{h}$.
To see this, note that for any $f\in\mathcal{F}_{h}$, there exists
$\tilde{f}\in\widetilde{\mathcal{F}}_{h}$ ($i=1,2,3$) for which
$\|M_{i}-\widetilde{M}_{i}\|_{\rm F}\leq\epsilon$, and thus

$\displaystyle\left\|\tilde{f}-f\right\|_{\infty}=$
$\displaystyle\max_{x,a,z}\Big{|}\mathcal{N}\big{(}z\,\big{|}\,x^{\top}M_{1}x+a^{\top}M_{2}x+a^{\top}M_{3}a,(H-h+1)\sigma^{2}\big{)}$
$\displaystyle-\mathcal{N}\big{(}z\,\big{|}\,x^{\top}\widetilde{M}_{1}x+a^{\top}\widetilde{M}_{2}x+a^{\top}\widetilde{M}_{3}a,(H-h+1)\sigma^{2}\big{)}\Big{|}$
$\displaystyle\leq$ $\displaystyle\frac{1}{(H-h+1)\sigma^{2}\sqrt{2\pi
e}}\underbrace{\left|x^{\top}(M_{1}-\widetilde{M}_{1})x+a^{\top}(M_{2}-\widetilde{M}_{2})x+a^{\top}(M_{3}-\widetilde{M}_{3})a\right|}_{(\heartsuit)}.$

where the last inequality holds for Lemma A.4. For $\rm(\heartsuit)$, we have

$\displaystyle(\heartsuit)\leq\|x\|_{2}\|M_{1}-\widetilde{M}_{1}\|_{\rm
F}\|x\|_{2}+\|a\|_{2}\|M_{2}-\widetilde{M}_{2}\|_{\rm
F}\|x\|_{2}+\|a\|_{2}\|M_{3}-\widetilde{M}_{3}\|_{\rm
F}\|a\|_{2}.\leq\epsilon(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2}).$

Hence, we have

$\left\|\tilde{f}-f\right\|_{\infty}\leq\epsilon\cdot\frac{m_{x}^{2}+m_{x}m_{a}+m_{a}^{2}}{(H-h+1)\sigma^{2}\sqrt{2\pi
e}}.$

This implies

$N\left(\frac{\epsilon(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{(H-h+1)\sigma^{2}\sqrt{2\pi
e}},\mathcal{F}_{h},\|\cdot\|_{\infty}\right)\leq
N(\epsilon,\mathcal{M}_{1},\|\cdot\|_{\rm F})\cdot
N(\epsilon,\mathcal{M}_{2},\|\cdot\|_{\rm F})\cdot
N(\epsilon,\mathcal{M}_{3},\|\cdot\|_{\rm F}).$

We note that $N_{[]}(2\epsilon,\mathcal{F}_{h},\|\cdot\|_{\infty})\leq
N(\epsilon,\mathcal{F}_{h},\|\cdot\|_{\infty})$. Hence we complete the proof.
∎

###### Proof of Lemma 4.17.

Let $\mathcal{M}_{i}=\\{M:\|M\|_{\rm F}\leq m_{i}\\}$ ($i=1,2,3$) denote the
set of possible matrices $M_{i}$. Then we have
$N(\epsilon,\mathcal{M}_{1},\|\cdot\|_{\rm
F})\leq(3m_{1}/\epsilon)^{d_{x}\times d_{x}}$,
$N(\epsilon,\mathcal{M}_{2},\|\cdot\|_{\rm
F})\leq(3m_{2}/\epsilon)^{d_{x}\times d_{a}}$, and
$N(\epsilon,\mathcal{M}_{3},\|\cdot\|_{\rm
F})\leq(3m_{3}/\epsilon)^{d_{a}\times d_{a}}$. By Lemma A.5, we have that

$\displaystyle N_{[]}(\epsilon,\mathcal{F}_{h},\|\cdot\|_{\infty})$
$\displaystyle\leq$
$\displaystyle\left(\frac{6m_{1}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}(H-h+1)\sqrt{2\pi
e}}\right)^{d_{x}\times
d_{x}}\left(\frac{6m_{2}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}(H-h+1)\sqrt{2\pi
e}}\right)^{d_{x}\times
d_{a}}\left(\frac{6m_{3}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}(H-h+1)\sqrt{2\pi
e}}\right)^{d_{a}\times d_{a}}$ $\displaystyle\leq$
$\displaystyle\left(\frac{6m_{1}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}\sqrt{2\pi
e}}\right)^{d_{x}\times
d_{x}}\left(\frac{6m_{2}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}\sqrt{2\pi
e}}\right)^{d_{x}\times
d_{a}}\left(\frac{6m_{3}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}\sqrt{2\pi
e}}\right)^{d_{a}\times d_{a}}.$

Taking a logarithm on both sides, we get

$\displaystyle\log N_{[]}($
$\displaystyle\epsilon,\mathcal{F}_{h},\|\cdot\|_{\infty})$
$\displaystyle\leq$ $\displaystyle
O\left(d_{x}^{2}\log\frac{m_{1}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}}+d_{x}d_{a}\log\frac{m_{2}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}}+d_{a}^{2}\log\frac{m_{3}(m_{x}^{2}+m_{x}m_{a}+m_{a}^{2})}{\epsilon\sigma^{2}}\right).$

∎

##  Appendix B Offline CVaR Evaluation

We consider estimating the CVaR of $Z^{\pi}$ with $d=1$. Given a threshold
$\tau\in(0,1)$, the $\text{CVaR}_{\tau}$ of $Z^{\pi}$ is defined as (assuming
finite-horizon MDPs):

$\displaystyle\text{CVaR}_{\tau}(Z^{\pi}):=\max_{b\in[0,H]}\left(b-\frac{1}{\tau}\mathbb{E}_{z\sim
Z^{\pi}}\max\left\\{b-z,0\right\\}\right).$

CVaR intuitively measures the expected value of the random variable belonging
to the tail part of the distribution and is often used as a risk-sensitive
measure. The following lemma shows that $\text{CVaR}_{\tau}(Z^{\pi})$ is a
Lipschitz continuous with respect to metric $d_{tv}$ and the Lipschitz
constant is $2H/\tau$.

######  Lemma B.1.

Let $f,f^{\prime}\in\Delta([0,H])$ be two densities. Then we have

$\displaystyle\text{\rm CVaR}_{\tau}(f)-\text{\rm
CVaR}_{\tau}(f^{\prime})\leq\frac{2H}{\tau}\cdot d_{tv}(f,f^{\prime}).$

###### Proof.

Let $f,f^{\prime}\in\Delta([0,H])$ denote two densities. Then we have

$\displaystyle\text{CVaR}_{\tau}(f)-\text{CVaR}_{\tau}(f^{\prime})$
$\displaystyle=$
$\displaystyle\max_{b\in[0,H]}\left(b-\frac{1}{\tau}\mathbb{E}_{z\sim
f}\max\left\\{b-z,0\right\\}\right)-\max_{b\in[0,H]}\left(b-\frac{1}{\tau}\mathbb{E}_{z\sim
f^{\prime}}\max\left\\{b-z,0\right\\}\right)$ $\displaystyle\leq$
$\displaystyle\left(b_{0}-\frac{1}{\tau}\mathbb{E}_{z\sim
f}\max\left\\{b_{0}-z,0\right\\}\right)-\left(b_{0}-\frac{1}{\tau}\mathbb{E}_{z\sim
f^{\prime}}\max\left\\{b_{0}-z,0\right\\}\right)$ $\displaystyle=$
$\displaystyle\frac{1}{\tau}\Big{(}\mathbb{E}_{z\sim
f^{\prime}}\max\left\\{b_{0}-z,0\right\\}-\mathbb{E}_{z\sim
f}\max\left\\{b_{0}-z,0\right\\}\Big{)}$ $\displaystyle=$
$\displaystyle\frac{1}{\tau}\int_{[0,H]}\big{(}f^{\prime}(z)-f(z)\big{)}\max\\{b_{0}-z,0\\}\mathop{}\\!\mathrm{d}z$
$\displaystyle\leq$
$\displaystyle\frac{H}{\tau}\int_{[0,H]}\big{|}f^{\prime}(z)-f(z)\big{|}\mathop{}\\!\mathrm{d}z$
$\displaystyle\leq$ $\displaystyle\frac{2H}{\tau}d_{tv}(f,f^{\prime})$

where the first inequality holds by picking
$b_{0}=\mathop{\arg\max}_{b\in[0,H]}\left(b-\frac{1}{\tau}\mathbb{E}_{z\sim
f}\max\left\\{b-z,0\right\\}\right)$. ∎

Thus using our bound from Corollary 4.7, we get:

$\displaystyle\left\lvert\text{CVaR}_{\tau}(Z^{\pi})-\text{CVaR}_{\tau}(\hat{f})\right\rvert\leq\frac{4C^{1/2}H^{2.5}}{\tau}\sqrt{\frac{\log(\max_{h}|\mathcal{F}|_{h}/\delta)}{n}},$

with probability at least $1-\delta$.

##  Appendix C Supporting Lemmas

###  C.1 Maximum Likelihood Estimation

In this section, we adapt the theoretical results of MLE [Agarwal et al.,
2020b] to more general versions. We will follow the notation in Appendix E of
Agarwal et al. [2020b] and restate the setting here for completeness.

We consider a sequential conditional probability estimation problem. Let
$\mathcal{X}$ and $\mathcal{Y}$ denote the instance space and the target
space, respectively. We are given a function class
$\mathcal{F}:(\mathcal{X}\times\mathcal{Y})\rightarrow\mathbb{R}$ with which
we want to model the true conditional distribution $f^{\star}$. To this end,
we are given a dataset
$D:=\left\\{\left(x_{i},y_{i}\right)\right\\}_{i=1}^{n}$, where
$x_{i}\sim\mathcal{D}_{i}$ and $y_{i}\sim p\left(\cdot\mid
x_{i}\right)=f^{\star}(x,\cdot)$.

We only assume that there exists $f_{i}^{\star}$ for each $i\in[n]$ such that
$\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}d_{tv}(f^{\star}_{i}(x),f^{\star}(x))=0$.
Note that this assumption only considers $x$ on the support of
$\mathcal{D}_{i}$ and is thus weaker than saying $f^{\star}\in\mathcal{F}$.

For the data generating process, we assume the data distribution
$\mathcal{D}_{i}$ is history-dependent, i.e., it can depend on the previous
samples: $x_{1},y_{1},\dots,x_{i-1},y_{i-1}$.

Let $\mathcal{D}^{\prime}=\\{(x_{i}^{\prime},y_{i}^{\prime})\\}_{i=1}^{n}$
denote the tangent sequence which is generated by
$x_{i}^{\prime}\sim\mathcal{D}_{i}$ and $y_{i}^{\prime}\sim p(\cdot\mid
x^{\prime}_{i})$. The tangent sequence is independent when conditioned on
$\mathcal{D}$.

######  Lemma C.1 (Adapted version of Lemma 25 [Agarwal et al., 2020b]).

Let $f_{1}\in\mathcal{X}\mapsto\Delta(\mathcal{Y})$ be a conditional
probability density and
$f_{2}\in\mathcal{X}\times\mathcal{Y}\mapsto\mathbb{R}_{\geq 0}$ (satisfying
$\int_{\mathcal{Y}}f_{2}(x,y)\mathop{}\\!\mathrm{d}y\leq s$ for all
$x\in\mathcal{X}$). Let $\mathcal{D}\in\Delta(\mathcal{X})$ be any
distribution. Then, we have

$\mathop{\mathbb{E}}_{x\sim\mathcal{D}}\left(\int_{\mathcal{Y}}\left|f_{1}(x,y)-f_{2}(x,y)\right|\mathop{}\\!\mathrm{d}y\right)^{2}\leq(2+2s)\left((s-1)-2\log\mathop{\mathbb{E}}_{x\sim\mathcal{D},y\sim
f_{1}(x,\cdot)}\exp\left(-\frac{1}{2}\log\left(f_{1}(x,y)/f_{2}(x,y)\right)\right)\right).$

###### Proof of Lemma C.1.

First, we have

$\displaystyle\mathop{\mathbb{E}}_{x\sim\mathcal{D}}\left(\int_{\mathcal{Y}}\left|f_{1}(x,y)-f_{2}(x,y)\right|\mathop{}\\!\mathrm{d}y\right)^{2}=\mathop{\mathbb{E}}_{x\sim\mathcal{D}}\left(\int_{\mathcal{Y}}\left|\sqrt{f_{1}(x,y)}-\sqrt{f_{2}(x,y)}\right|\left(\sqrt{f_{1}(x,y)}+\sqrt{f_{2}(x,y)}\right)\mathop{}\\!\mathrm{d}y\right)^{2}$
$\displaystyle\leq$
$\displaystyle\mathop{\mathbb{E}}_{x\sim\mathcal{D}}{\int_{\mathcal{Y}}\left(\sqrt{f_{1}(x,y)}-\sqrt{f_{2}(x,y)}\right)^{2}\mathop{}\\!\mathrm{d}y}\cdot{\int_{\mathcal{Y}}\left(\sqrt{f_{1}(x,y)}+\sqrt{f_{2}(x,y)}\right)^{2}\mathop{}\\!\mathrm{d}y}$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}_{x\sim\mathcal{D}}{\int_{\mathcal{Y}}\left(\sqrt{f_{1}(x,y)}-\sqrt{f_{2}(x,y)}\right)^{2}\mathop{}\\!\mathrm{d}y}\cdot
2\int_{\mathcal{Y}}\big{(}f_{1}(x,y)+f_{2}(x,y)\big{)}\mathop{}\\!\mathrm{d}y-\int_{\mathcal{Y}}\left(\sqrt{f_{1}(x,y)}-\sqrt{f_{2}(x,y)}\right)^{2}\mathop{}\\!\mathrm{d}y$
$\displaystyle=$
$\displaystyle\mathop{\mathbb{E}}_{x\sim\mathcal{D}}{\int_{\mathcal{Y}}\left(\sqrt{f_{1}(x,y)}-\sqrt{f_{2}(x,y)}\right)^{2}\mathop{}\\!\mathrm{d}y}\cdot
2\int_{\mathcal{Y}}\big{(}f_{1}(x,y)+f_{2}(x,y)\big{)}\mathop{}\\!\mathrm{d}y$
$\displaystyle\leq$
$\displaystyle\underbrace{\mathop{\mathbb{E}}_{x\sim\mathcal{D}}\int_{\mathcal{Y}}\left(\sqrt{f_{1}(x,y)}-\sqrt{f_{2}(x,y)}\right)^{2}\mathop{}\\!\mathrm{d}y}_{\rm(*)}\cdot(2+2s).$

where the first inequality holds for Cauchy–Schwarz inequality. For
${\rm(*)}$, we have

$\displaystyle{\rm(*)}=$
$\displaystyle\mathop{\mathbb{E}}_{x\sim\mathcal{D}}\int_{\mathcal{Y}}\left(\sqrt{f_{1}(x,y)}-\sqrt{f_{2}(x,y)}\right)^{2}\mathop{}\\!\mathrm{d}y\leq(s-1)+2-2\mathop{\mathbb{E}}_{x\sim\mathcal{D}}\int_{\mathcal{Y}}\sqrt{f_{1}(x,y)f_{2}(x,y)}\mathop{}\\!\mathrm{d}y$
$\displaystyle=$
$\displaystyle(s-1)+2\left(1-\mathop{\mathbb{E}}_{x\sim\mathcal{D}}\int_{\mathcal{Y}}\sqrt{f_{1}(x,y)f_{2}(x,y)}\mathop{}\\!\mathrm{d}y\right)\leq(s-1)-2\log\left(\mathop{\mathbb{E}}_{x\sim\mathcal{D}}\int_{\mathcal{Y}}\sqrt{f_{1}(x,y)f_{2}(x,y)}\mathop{}\\!\mathrm{d}y\right)$
$\displaystyle\leq$
$\displaystyle(s-1)-2\log\mathop{\mathbb{E}}_{x\sim\mathcal{D},y\sim
f_{1}(x,\cdot)}\sqrt{f_{2}(x,y)/f_{1}(x,y)}$ $\displaystyle=$
$\displaystyle(s-1)-2\log\mathop{\mathbb{E}}_{x\sim\mathcal{D},y\sim
f_{1}(x,\cdot)}\exp\left(-\frac{1}{2}\log\left(f_{1}(x,y)/f_{2}(x,y)\right)\right)$

where the second inequality holds because $1-x\leq-\log x$. ∎

######  Lemma C.2 (Adapted version of Theorem 21 [Agarwal et al., 2020b]).

Fix $\delta\in(0,1)$. Let $N_{[]}(\epsilon,\mathcal{F},\|\cdot\|_{\infty})$
denote the $\epsilon$-bracketing number of $\mathcal{F}$ w.r.t.
$\|\cdot\|_{\infty}$. Then for any estimator $\hat{f}$ that depends on $D$,
with probability at least $1-\delta$, we have

$\displaystyle\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}$
$\displaystyle d_{tv}^{2}\left(\hat{f}(x,\cdot),f^{\star}(x,\cdot)\right)\leq$
$\displaystyle\frac{3n\epsilon^{2}|\mathcal{Y}|^{2}}{2}+2n\epsilon|\mathcal{Y}|+\big{(}4+2\epsilon|\mathcal{Y}|\big{)}\left(\frac{1}{2}\sum_{i=1}^{n}\log\big{(}f^{\star}(x_{i},y_{i})/\hat{f}(x_{i},y_{i})\big{)}+\log
N_{[]}(\epsilon,\mathcal{F},\|\cdot\|_{\infty})+\log(1/\delta)\right)$

where $|\mathcal{Y}|$ denotes $\int_{\mathcal{Y}}\mathop{}\\!\mathrm{d}y$.

###### Proof of Lemma C.2.

We take an $\epsilon$-bracket of $\mathcal{F}$,
$\\{[l_{i},u_{i}]:i=1,2,\dots\\}$, and denote
$\widetilde{\mathcal{F}}=\\{u_{i}:i=1,2,\dots\\}$. Pick
$\tilde{f}\in\widetilde{\mathcal{F}}$ satisfying $\hat{f}\leq\tilde{f}$, so
$\tilde{f}$ also depends on $D$. Applying Lemma 24 of Agarwal et al. [2020b]
to function class $\widetilde{\mathcal{F}}$ and estimator $\tilde{f}$ and
using Chernoff method, we have

$\underbrace{-\log\mathop{\mathbb{E}}_{D^{\prime}}\exp(L(\tilde{f}(D),D^{\prime}))}_{\rm(i)}\leq\underbrace{-L(\tilde{f}(D),D)+\log
N_{[]}(\epsilon,\mathcal{F},\|\cdot\|_{\infty})+\log(1/\delta)}_{\rm(ii)}.$
(9)

holds with probability at least $1-\delta$. We set
$L(f,D)=\sum_{i=1}^{n}-\nicefrac{{1}}{{2}}\log(f^{\star}(x_{i},y_{i})/f(x_{i},y_{i}))$.
Then the right hand side of (9) is

$\displaystyle{\rm{(ii)}}=$
$\displaystyle\frac{1}{2}\sum_{i=1}^{n}\log(f^{\star}(x_{i},y_{i})/\tilde{f}(x_{i},y_{i}))+\log
N_{[]}(\epsilon,\mathcal{F},\|\cdot\|_{\infty})+\log(1/\delta)$
$\displaystyle\leq$
$\displaystyle\frac{1}{2}\sum_{i=1}^{n}\log(f^{\star}(x_{i},y_{i})/\hat{f}(x_{i},y_{i}))+\log
N_{[]}(\epsilon,\mathcal{F},\|\cdot\|_{\infty})+\log(1/\delta).$

On the other hand, by the definition of total variation distance and the fact
that $a^{2}\leq 2b^{2}+2c^{2}$ whenever $0\leq a\leq b+c$, we have

$\displaystyle\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}d_{tv}^{2}\left(\hat{f}(x,\cdot),f^{\star}(x,\cdot)\right)=\frac{1}{4}\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}\left(\int_{\mathcal{Y}}\left|\hat{f}(x,y)-f^{\star}(x,y)\right|\mathop{}\\!\mathrm{d}y\right)^{2}$
$\displaystyle\leq$
$\displaystyle\frac{1}{2}\underbrace{\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}\left(\int_{\mathcal{Y}}\left|\hat{f}(x,y)-\tilde{f}(x,y)\right|\mathop{}\\!\mathrm{d}y\right)^{2}}_{\rm{(iii)}}+\frac{1}{2}\underbrace{\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}\left(\int_{\mathcal{Y}}\left|\tilde{f}(x,y)-f^{\star}(x,y)\right|\mathop{}\\!\mathrm{d}y\right)^{2}}_{\rm{(iv)}}.$

For $\rm{(iii)}$, by the definition of $\tilde{f}$, we have ${\rm{(iii)}}\leq
n\epsilon^{2}|\mathcal{Y}|^{2}$. For $\rm{(iv)}$, we apply Lemma C.1 with
$f_{1}=f^{\star}$ and $f_{2}=\tilde{f}$ (thus $s=1+\epsilon|\mathcal{Y}|$) and
get

$\displaystyle{\rm(iv)}=$ $\displaystyle
2n\epsilon|\mathcal{Y}|(2+\epsilon|\mathcal{Y}|)-\sum_{i=1}^{n}(8+4\epsilon|\mathcal{Y}|)\left(\log\mathop{\mathbb{E}}_{x,y\sim
f^{\star}(x,\cdot)}\exp\left(-\frac{1}{2}\log\left(f^{\star}(x,y)/\tilde{f}(x,y)\right)\right)\right)$
$\displaystyle=$ $\displaystyle
2n\epsilon|\mathcal{Y}|(2+\epsilon|\mathcal{Y}|)-\sum_{i=1}^{n}(8+4\epsilon|\mathcal{Y}|)\left(\log\mathop{\mathbb{E}}_{x,y\sim\mathcal{D}_{i}}\exp\left(-\frac{1}{2}\log\left(f^{\star}(x,y)/\tilde{f}(x,y)\right)\right)\right)$
$\displaystyle=$ $\displaystyle
2n\epsilon|\mathcal{Y}|(2+\epsilon|\mathcal{Y}|)-(8+4\epsilon|\mathcal{Y}|)\log\mathop{\mathbb{E}}_{x,y\sim\mathcal{D}^{\prime}}\left[\exp\left(\sum_{i=1}^{n}-\frac{1}{2}\log\left(f^{\star}(x,y)/\tilde{f}(x,y)\right)\right)\middle|D\right]$
$\displaystyle=$ $\displaystyle
4n\epsilon|\mathcal{Y}|+2n\epsilon^{2}|\mathcal{Y}|^{2}+(8+4\epsilon|\mathcal{Y}|)\cdot{\rm{(i)}}.$

By plugging $\rm{(iii)}$ and $\rm{(iv)}$ back we get

$\displaystyle\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}d_{tv}^{2}\left(\hat{f}(x,\cdot),f^{\star}(x,\cdot)\right)\leq
2n\epsilon|\mathcal{Y}|+\frac{3}{2}n\epsilon^{2}|\mathcal{Y}|^{2}+(4+2\epsilon|\mathcal{Y}|)\cdot{\rm{(i)}}.$

Notice that $\rm{(i)}\leq\rm{(ii)}$, so we complete the proof by plugging
$\rm(ii)$ into the above. ∎

######  Lemma C.3.

Fixed $\delta\in(0,1)$. Let $\hat{f}$ denote the maximum likelihood estimator,

$\hat{f}=\mathop{\arg\max}_{f\in\mathcal{F}}\sum_{i=1}^{n}\log
f(x_{i},y_{i}).$

Then according to different assumptions on the size of $\mathcal{F}$, we have
the following two conclusions:

  1. (1)

If $|\mathcal{F}|<\infty$, we have

$\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}d_{tv}^{2}\left(\hat{f}(x,\cdot),f^{\star}(x,\cdot)\right)\leq
4\log|\mathcal{F}|/\delta$ (10)

with probability at least $1-\delta$.

  2. (2)

For general $\mathcal{F}$, we have

$\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}d_{tv}^{2}\left(\hat{f}(x,\cdot),f^{\star}(x,\cdot)\right)\leq
10\log
N_{[]}\left((n|\mathcal{Y}|)^{-1},\mathcal{F},\|\cdot\|_{\infty}\right)/\delta$
(11)

with probability at least $1-\delta$.

###### Proof of Lemma C.3.

By Lemma C.2, we have

$\displaystyle\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}$
$\displaystyle d_{tv}^{2}\left(\hat{f}(x,\cdot),f^{\star}(x,\cdot)\right)\leq$
(12)
$\displaystyle\frac{3n\epsilon^{2}|\mathcal{Y}|^{2}}{2}+2n\epsilon|\mathcal{Y}|+\big{(}4+2\epsilon|\mathcal{Y}|\big{)}\left(\frac{1}{2}\underbrace{\sum_{i=1}^{n}\log\big{(}f^{\star}(x_{i},y_{i})/\hat{f}(x_{i},y_{i})\big{)}}_{\rm(\diamond)}+\log
N_{[]}(\epsilon,\mathcal{F},\|\cdot\|_{\infty})+\log(1/\delta)\right)$

with probability at least $1-\delta$. Since $\hat{f}$ is the maximum
likelihood estimator and there exists $f^{\star}_{i}$ that agrees with
$f^{\star}$ on the support of $\mathcal{D}_{i}$, we have

$\displaystyle\log\big{(}f^{\star}(x_{i},y_{i})/\hat{f}(x_{i},y_{i})\big{)}=\log\big{(}f^{\star}_{i}(x_{i},y_{i})/\hat{f}(x_{i},y_{i})\big{)}\leq
0$

and thus${\rm(\diamond)}\leq 0$. When $|\mathcal{F}|<\infty$, we can set
$\epsilon=0$, and then (12) exactly becomes (10). For general $\mathcal{F}$,
we set $\epsilon=(n|\mathcal{Y}|)^{-1}$ and then get

$\displaystyle\sum_{i=1}^{n}\mathop{\mathbb{E}}_{x\sim\mathcal{D}_{i}}d_{tv}^{2}\left(\hat{f}(x,\cdot),f^{\star}(x,\cdot)\right)\leq\frac{3}{2n}+2+\left(4+\frac{2}{n}\right)\log
N_{[]}((n|\mathcal{Y}|)^{-1},\mathcal{F},\|\cdot\|_{\infty})/\delta)$
$\displaystyle\leq$ $\displaystyle 4+6\log
N_{[]}((n|\mathcal{Y}|)^{-1},\mathcal{F},\|\cdot\|_{\infty})/\delta\leq 10\log
N_{[]}((n|\mathcal{Y}|)^{-1},\mathcal{F},\|\cdot\|_{\infty})/\delta,$

which is exactly (11). ∎

###  C.2 Total Variation Distance and Wasserstein Distance

The following lemma states that the total variation distance is equal to the
optimal coupling in a sense. The proof can be found in Levin and Peres [2017]
(Proposition 4.7).

######  Lemma C.4.

Let $f_{1}$ and $f_{2}$ be two probability distributions on $\mathcal{X}$.
Then

$\displaystyle d_{tv}(f_{1},f_{2})=\inf_{c\in\mathcal{C}}\Pr_{x,y\sim c}(x\neq
y)$

where $\mathcal{C}$ is the set of all couplings of $f_{1}$ and $f_{2}$.

The following lemma shows the dual representation of the Wasserstein distance.
The proof can be found in Villani [2021] (Theorem 1.3) and Villani et al.
[2009] (Theorem 5.10).

######  Lemma C.5 (Kantorovich duality).

Let $f_{1},f_{2}\in\Delta(\mathcal{X})$ where $\mathcal{X}$ is a Polish space
(e.g., Euclidean space). It can be shown that, for any $1\leq p<\infty$,

$\displaystyle
d_{w,p}^{p}(f_{1},f_{2})=\sup_{\psi,\phi}\int\psi(x)f_{1}(x)\mathop{}\\!\mathrm{d}x-\int\phi(x)f_{2}(x)\mathop{}\\!\mathrm{d}x\text{\quad
s.t.\quad}\psi(x)-\phi(y)\leq\|x-y\|^{p},\quad\forall x,y\in\mathcal{X}.$

######  Lemma C.6.

Let $f_{1}$ and $f_{2}$ be two distributions on a bounded set $\mathcal{X}$.
Then

$d_{w,p}^{p}(f_{1},f_{2})\leq\text{\rm diam}^{p}(\mathcal{X})\cdot
d_{tv}(f_{1},f_{2})$

where $\text{\rm diam($\mathcal{X}$)}=\sup_{x,y\in\mathcal{X}}\|x-y\|$ is the
diameter of $\mathcal{X}$.

###### Proof.

By definition, we have

$\displaystyle
d_{w,p}^{p}(f_{1},f_{2})=\inf_{c\in\mathcal{C}}\mathop{\mathbb{E}}_{x,y\sim
c}\|x-y\|^{p}=\inf_{c\in\mathcal{C}}\mathop{\mathbb{E}}_{x,y\sim
c}\big{[}\mathds{1}[x\neq y]\cdot\|x-y\|^{p}\big{]}$ $\displaystyle\leq$
$\displaystyle\text{\rm
diam}^{p}(\mathcal{X})\cdot\inf_{c\in\mathcal{C}}\mathop{\mathbb{E}}_{x,y\sim
c}\mathds{1}[x\neq y]=\text{\rm diam}^{p}(\mathcal{X})\cdot
d_{tv}(f_{1},f_{2})$

where by $\mathcal{C}$ we denote the set of all couplings of $f_{1}$ and
$f_{2}$, and the last equality holds because of Lemma C.4. ∎

######  Corollary C.7.

Let $f_{1}$ and $f_{2}$ be two distributions on $[0,m]^{d}$. Then

$d_{w,p}^{p}(f_{1},f_{2})\leq\left(m\sqrt{d}\right)^{p}\cdot
d_{tv}(f_{1},f_{2}).$

Since the total variation distance is at most one, we have the following.

######  Corollary C.8.

Let $f_{1}$ and $f_{2}$ be two distributions on $[0,m]^{d}$. Then

$d_{w,p}^{p}(f_{1},f_{2})\leq\left(m\sqrt{d}\right)^{p}.$

##  Appendix D Experiment Details

###  D.1 Implementation Details of Combination Lock Environment

We first clarify our implementation of the combination lock environment.

##### Reward.

We denote $r^{+}$ and $r^{-}$ as the random reward for latent state $w_{H}=0$
and $w_{H}=1$, respectively. For the one-dimensional case, they are sampled
from Gaussian distributions: $r^{+}\sim\mathcal{N}(1,0.1^{2})$ and
$r^{-}\sim\mathcal{N}(-1,0.1^{2})$. For the second experiment with two-
dimensional reward, they are defined as

$\displaystyle r^{+}=x+\frac{2x}{\|x\|_{2}}\quad\text{where}\quad
x\sim\mathcal{N}\left(\begin{bmatrix}0\\\ 0\end{bmatrix},\
\begin{bmatrix}0.05&0\\\ 0&0.05\end{bmatrix}\right),\quad
r^{-}\sim\mathcal{N}\left(\begin{bmatrix}0\\\ 0\end{bmatrix},\
\begin{bmatrix}0.05&0\\\ 0&0.05\end{bmatrix}\right).$

Visually, most samples of $r^{+}$ appear in a ring centered at the origin with
a radius of 2.

##### State.

The state is constructed by three components, that is, state
$x=(x_{1},x_{2},x_{3})^{\top}$ for which $x_{1}$ is the one-hot encoding of
latent state, $x_{2}$ is the one-hot encoding of time step $h$, and $x_{3}$ is
a vector of Gaussian noise sampled independently from
$\mathcal{N}(0,0.1^{2})$.

The optimal action $a^{\star}_{h}$ is chosen to be 0 for all $h\in[H]$ for
simplicity. We list other environment hyperparameters in Table 3 for
reference.

Table 3: Hyperparameters for the combination lock environment. The two columns
denote the respective hyperparameters employed in one-dimensional and two-
dimensional experiments.
| 1-dimensional | 2-dimensional
---|---|---
Horizon | 20 | 10
Number of Actions | 2 | 2
Dimension of States | 30 | 30

###  D.2 Implementation Details of Algorithms

All algorithms, with the exception of Diff-FLE, is implemented by a neural
network consisting of two layers, each with 32 neurons, connected by the ReLU
activation functions. Diff-FLE employs a three-layered neural network, each
layer containing 256 neurons, connected by the ReLU functions. Some shared
hyperparameters are listed in Table 4.

Table 4: Shared hyperparameters. Note that the size of the dataset is written
as a product, which is determined by the way we generate the offline data: the
first number means the number of samples generated for each latent state and
each time step, the second number means the number of time steps (i.e.,
horizon), and the third number means the size of the latent space.
| 1-dimensional | 2-dimensional
---|---|---
Size of Dataset | $10000\times 20\times 2$ | $10000\times 10\times 2$
Batch Size | 500 | 500

##### Categorical Algorithm.

We present the implementation of the two-dimensional version of the
categorical algorithm, which is not presented in the prior work [Bellemare et
al., 2017]. As a reminder, for the one-dimensional counterpart, for each atom
of the next state, we first calculate its target position, then distribute the
probability of that atom based on the distance of the target position to the
closest two atoms. In the two-dimensional case, we discretize on each
dimension, resulting in a grid-shaped discretization. Therefore, the
probability of the atoms of the next state will be distributed based on the
distance to the _four_ closest atoms (generally, it will be distributed to
$2^{n}$ atoms in the $n$-dimensional case). The other implementation details
are the same as the one-dimensional case. The list of hyperparameters can be
found in the Table 5.

Table 5: Hyperparameters for the categorical algorithm.
| 1-dimensional | 2-dimensional
---|---|---
Number of Atoms | $100$ | $30^{2}$
Learning Rate | $10^{-2}$ | $3\times 10^{-2}$
Number of Iterations | 200 | 100
Discretized Range | $[-1.5,1.5]$ | $[-4,4]^{2}$

##### Quantile Algorithm.

We followed the implementation of Dabney et al. [2018]. The list of
hyperparameters can be found in the Table 6.

Table 6: Hyperparameters for quantile Algorithm.
| 1-dimensional
---|---
Number of Quantiles | 100
Learning Rate | $10^{-3}$
Number of Iterations | 1000

##### Diff-FLE.

Our implementation is based on DDPM [Ho et al., 2020]. However, our neural
network is much simpler from theirs, as mentioned above. The list of
hyperparameters can be found in the Table 7.

Table 7: Hyperparameters for Diff-FLE.
| 1-dimensional | 2-dimensional
---|---|---
Steps of Diffusion Process | 200 | 200
Staring Variance | $10^{-3}$ | $10^{-3}$
Final Variance | 0.1 | 0.1
Variance Increasing | Linear | Linear
Learning Rate | $10^{-3}$ | $10^{-3}$
Number of Iterations | 5000 | 15000

##### GMM-FLE.

For the training of GMM-FLE, we applied gradient ascent on the log-likelihood.
While many classic approaches (e.g., the Expectation-Maximization (EM)
algorithm) exist, we found no significant performance gap between gradient
ascent and EM in our trials on both one-dimensional and two-dimensional data.
Therefore, we opted for the gradient ascent, which matches our theory better.
The list of hyperparameters is listed in Table 8.

Table 8: Hyperparameters for GMM-FLE.
| 1-dimensional | 2-dimensional
---|---|---
Number of Gaussian Distribution | 10 | 10
Learning Rate | $10^{-4}$ | $2\times 10^{-4}$
Number of Iterations | 20000 | 10000

###  D.3 Full Experiment Results

Table 9 is the full version of Table 1, and Table 10 is the full version of
Table 2.

Table 9: Full version of Table 1.
$h$ | Cate Alg | Quan Alg | Diff-FLE | GMM-FLE
---|---|---|---|---
1 |  0.071 $\pm$ 0.015 |  0.603 $\pm$ 0.011 |  0.292 $\pm$ 0.073 |  0.039 $\pm$ 0.004
2 |  0.067 $\pm$ 0.012 |  0.609 $\pm$ 0.014 |  0.305 $\pm$ 0.055 |  0.041 $\pm$ 0.005
3 |  0.068 $\pm$ 0.013 |  0.612 $\pm$ 0.017 |  0.305 $\pm$ 0.079 |  0.039 $\pm$ 0.009
4 |  0.073 $\pm$ 0.013 |  0.593 $\pm$ 0.015 |  0.288 $\pm$ 0.073 |  0.038 $\pm$ 0.003
5 |  0.074 $\pm$ 0.015 |  0.602 $\pm$ 0.009 |  0.285 $\pm$ 0.054 |  0.036 $\pm$ 0.009
6 |  0.077 $\pm$ 0.011 |  0.612 $\pm$ 0.010 |  0.268 $\pm$ 0.040 |  0.030 $\pm$ 0.008
7 |  0.080 $\pm$ 0.014 |  0.602 $\pm$ 0.014 |  0.290 $\pm$ 0.066 |  0.034 $\pm$ 0.004
8 |  0.080 $\pm$ 0.016 |  0.584 $\pm$ 0.018 |  0.273 $\pm$ 0.039 |  0.039 $\pm$ 0.013
9 |  0.081 $\pm$ 0.019 |  0.529 $\pm$ 0.028 |  0.247 $\pm$ 0.034 |  0.048 $\pm$ 0.010
10 |  0.079 $\pm$ 0.017 |  0.494 $\pm$ 0.018 |  0.234 $\pm$ 0.043 |  0.044 $\pm$ 0.012
11 |  0.080 $\pm$ 0.016 |  0.514 $\pm$ 0.018 |  0.244 $\pm$ 0.038 |  0.039 $\pm$ 0.012
12 |  0.089 $\pm$ 0.009 |  0.518 $\pm$ 0.013 |  0.232 $\pm$ 0.015 |  0.032 $\pm$ 0.007
13 |  0.089 $\pm$ 0.011 |  0.481 $\pm$ 0.016 |  0.219 $\pm$ 0.027 |  0.029 $\pm$ 0.016
14 |  0.081 $\pm$ 0.015 |  0.416 $\pm$ 0.026 |  0.221 $\pm$ 0.021 |  0.033 $\pm$ 0.012
15 |  0.083 $\pm$ 0.015 |  0.330 $\pm$ 0.028 |  0.178 $\pm$ 0.033 |  0.026 $\pm$ 0.015
16 |  0.081 $\pm$ 0.009 |  0.283 $\pm$ 0.017 |  0.170 $\pm$ 0.045 |  0.027 $\pm$ 0.013
17 |  0.082 $\pm$ 0.008 |  0.252 $\pm$ 0.008 |  0.167 $\pm$ 0.037 |  0.034 $\pm$ 0.013
18 |  0.070 $\pm$ 0.010 |  0.217 $\pm$ 0.012 |  0.133 $\pm$ 0.019 |  0.023 $\pm$ 0.008
19 |  0.078 $\pm$ 0.011 |  0.167 $\pm$ 0.019 |  0.109 $\pm$ 0.031 |  0.018 $\pm$ 0.008
20 |  0.077 $\pm$ 0.014 |  0.076 $\pm$ 0.009 |  0.067 $\pm$ 0.024 |  0.013 $\pm$ 0.005
Table 10: Full version of Table 2.
$h$ | Cate Alg | Diff-FLE | GMM-FLE
---|---|---|---
1 |  0.483 $\pm$ 0.003 |  0.357 $\pm$ 0.031 |  0.438 $\pm$ 0.008
2 |  0.483 $\pm$ 0.003 |  0.344 $\pm$ 0.030 |  0.424 $\pm$ 0.048
3 |  0.480 $\pm$ 0.003 |  0.339 $\pm$ 0.023 |  0.450 $\pm$ 0.042
4 |  0.469 $\pm$ 0.002 |  0.327 $\pm$ 0.019 |  0.478 $\pm$ 0.048
5 |  0.466 $\pm$ 0.001 |  0.310 $\pm$ 0.019 |  0.493 $\pm$ 0.050
6 |  0.466 $\pm$ 0.001 |  0.289 $\pm$ 0.031 |  0.491 $\pm$ 0.061
7 |  0.470 $\pm$ 0.003 |  0.256 $\pm$ 0.032 |  0.510 $\pm$ 0.080
8 |  0.465 $\pm$ 0.002 |  0.234 $\pm$ 0.023 |  0.505 $\pm$ 0.099
9 |  0.453 $\pm$ 0.001 |  0.207 $\pm$ 0.014 |  0.502 $\pm$ 0.094
10 |  0.446 $\pm$ 0.002 |  0.143 $\pm$ 0.011 |  0.376 $\pm$ 0.101