SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alexandre Marthe
UMPA
ENS de Lyon
alexandre.marthe@ens-lyon.fr
Aurelien Garivier
UMPA UMR 5669 and LIP UMR 5668
Univ. Lyon, ENS de Lyon
46 allée d'Italie F-69364 Lyon cedex 07, France
aurelien.garivier@ens-lyon.fr
Claire Vernade
University of Tuebingen
Tuebingen, Germany
claire.vernade@uni-tuebingen.de
What are the functionals of the reward that can be computed and optimized exactly in Markov Decision Processes? In the finite-horizon, undiscounted setting, Dynamic Programming (DP) can only handle these operations efficiently for certain classes of statistics. We summarize the characterization of these classes for policy evaluation, and give a new answer for the planning problem. Interestingly, we prove that only generalized means can be optimized exactly, even in the more general framework of Distributional Reinforcement Learning (DistRL). DistRL permits, however, to evaluate other functionals approximately. We provide error bounds on the resulting estimators, and discuss the potential of this approach as well as its limitations. These results contribute to advancing the theory of Markov Decision Processes by examining overall characteristics of the return, and particularly risk-conscious strategies.
Reinforcement Learning (RL) has emerged as a flourishing field of study, delivering significant practical applications ranging from robot control and game solving to drug discovery or hardware design [Lazic et al., 2018, Popova et al., 2018, Volk et al., 2023, Mirhoseini et al., 2020]. The cornerstone of RL is the "return" value, a sum of successive rewards. Conventionally, the focus is on computing and optimizing its expected value on Markov Decision Process (MDP). The remarkable efficiency of MDPs comes from their ability to be solved through dynamic programming with the Bellman equations [Sutton and Barto, 2018, Szepesvári, 2010]. RL theory has seen considerable expansion, with a renewed interest for the consideration of more rich descriptions of a policy's behavior than the sole average return. At the other end of the spectrum, the so-called Distributional Reinforcement Learning (DistRL) approach aims at studying and optimizing the entire return distribution, leading to impressive practical results [Bellemare et al., 2017, Hessel et al., 2018, Wurman et al., 2022, Fawzi et al., 2022]. Between the expectation and the entire distribution, the efficient handling of other statistical functionals of the reward appears also particularly relevant for risk-sensitive contexts [Bernhard et al., 2019, Mowbray et al., 2022].
Despite recent progress, the full understanding of the abilities and limitations of DistRL to compute other functionals remains incomplete, with the underlying theory yet to be fully understood. Historically, the theory of RL has been established for discounted MDPs, see e.g. [Sutton and Barto, 2018, Watkins and Dayan, 1992, Szepesvári, 2010, Bellemare et al., 2023] for modern reference textbooks. Recently more attention was drawn to the undiscounted, finite-horizon setting [Auer, 2002,
Osband et al., 2013, Jin et al., 2018, Ghavamzadeh et al., 2020], for which fundamental questions remain open. In this paper, we explore policy evaluation, planning and exact learning algorithms for undiscounted MDPs for the optimization problem of general functionals of the reward. We explicitly delimit the possibilities offered by dynamic programming as well as DistRL.
Our paper specifically addresses two questions:
(i) How accurately can we evaluate statistical functionals by using DistRL?
(ii) Which functionals can be exactly optimized through dynamic programming?
We first recall the fundamental results in dynamic programming and Distributional RL. Addressing question (i), we refer to Rowland et al. [2019]'s results on Bellman closedness and provide their adaptation to undiscounted MDPs. We then prove upper bounds on the approximation error of Policy Evaluation with DistRL and corroborate these bounds with practical experiments. For question (ii), we draw a connection between Bellman closedness and planning. We then utilize the DistRL framework to identify two key properties held by estimizable functionals.
Our main contribution is a characterization of the families of utilities that verify these two properties (Theorem 2). This result gives a comprehensive answer to question (ii) and closes an important open issue in the theory of MDP. It shows in particular that DistRL does not extend the class of functionals for which planing is possible beyond what is already allowed by classical dynamic programming.
We introduce the classical RL framework in finite-horizon tabular Markov Decisions Processes (MDPs). We write  $\mathcal{P}(\mathbb{R})$  the space of probability distributions on  $\mathbb{R}$ . A finite-horizon tabular MDP is a tuple  $\mathcal{M} = (\mathcal{X},\mathcal{A},p,R,H)$ , where  $\mathcal{X}$  is a finite state space,  $\mathcal{A}$  is a finite action space,  $H$  is the horizon, for each  $h\in [H]$ ,  $p_h(x,a,\cdot)$  is a transition probability law and  $R_{h}(x,a)$  is a reward random variable with distribution  $\varrho_{h}$ . The parameters  $(p_h)$  and  $(R_{h})$  define the model of dynamics. A deterministic policy on  $\mathcal{M}$  is a sequence  $\pi = (\pi_1,\dots ,\pi_H)$  of functions  $\pi_h:\mathcal{X}\to \mathcal{A}$ .
Reinforcement Learning traditionally focuses on learning policies optimizing the expected return. For a given policy  $\pi$ , the  $Q$ -function maps a state-action pair to its expected return under  $\pi$ :
$$
Q _ {h} ^ {\pi} (x, a) = \mathbb {E} _ {\varrho_ {h}} \left[ R _ {h} (x, a) \right] + \sum_ {s ^ {\prime}} p _ {h} (x, a, x ^ {\prime}) Q _ {h + 1} ^ {\pi} \left(x ^ {\prime}, \pi_ {h + 1} \left(x ^ {\prime}\right)\right), \quad Q _ {H + 1} ^ {\pi} (x, a) = 0. \tag {1}
$$
When the model is known, the  $Q$ -function of a policy  $\pi$  can be computed by doing a backward recursion, also called dynamic programming. This is referred to as Policy Evaluation. Similarly, an optimal policy can be found by solving the optimal Bellman equation:
$$
Q _ {h} ^ {*} (x, a) = \mathbb {E} _ {\varrho_ {h}} \left[ R _ {h} (x, a) \right] + \sum_ {x ^ {\prime}} p _ {h} (x, a, x ^ {\prime}) \max  _ {a ^ {\prime}} Q _ {h + 1} ^ {*} \left(x ^ {\prime}, a ^ {\prime}\right), \quad Q _ {H + 1} ^ {*} (x, a) = 0. \tag {2}
$$
Solving this equation when the model is known is also called Planning. When it is unknown, reinforcement learning aims at finding the optimal policy from sample runs of the MDP. But evaluating and optimizing the expectation of the return in the definition of the Q-function above is just one choice of statistical functional. We now introduce Distributional RL and then discuss other statistical functionals that generalize the expected setting discussed so far.
2.1 Distributional RL
Distributional RL (DistRL) refers to the approach that tracks not just a statistic of the return for each state but its entire distribution. We introduce here the most important basic concepts and refer the reader to the recent comprehensive survey by Bellemare et al. [2023] for more details. The main idea is to use the full distributions to estimate and optimize various metrics over the returns ranging from the mere expectation [Bellemare et al., 2017] to more complex metrics [Rowland et al., 2019, Dabney et al., 2018a, Liang and Luo, 2022].
At state-action  $(x, a)$ , let  $Z_h^\pi(x, a)$  denote the future sum of rewards when following policy  $\pi$  and starting at step  $h$ , also called return. It verifies the simple recursive formula  $Z_h^\pi(x, a) = R_h(x, a) + Z_{h+1}^\pi(X', \pi_{h+1}(X'))$  where  $X' \sim p_{h+1}(x, a, \cdot)$ . Its distribution is  $\eta = (\eta_h^\pi(x, a))_{(x, a, h) \in \mathcal{X} \times \mathcal{A} \times [H]}$  and is often referred to as the  $Q$ -value distribution. One can easily derive the recursive law of the return as a convolution: for any two measures  $\nu_1, \nu_2 \in \mathcal{P}(\mathbb{R})$ , we denote their convolution by  $\nu_1 * \nu_2(t) = \int_{\mathbb{R}} \nu_1(\tau) \nu_2(t - \tau) d\tau$ . For any two independent random variables  $X$  and  $Y$ , the distribution of the sum  $Z = X + Y$  is the convolution of their distributions:  $\nu_Z = \nu_X * \nu_Y$ . Thus, the law of  $Z_h^\pi(x, a)$  is
$$
\forall x, a, h, \quad \eta_ {h} ^ {\pi} (x, a) = \varrho_ {h} (x, a) * \sum_ {x ^ {\prime}} p _ {h} \left(x, a, x ^ {\prime}\right) \eta_ {h} ^ {\pi} \left(x ^ {\prime}, \pi_ {h + 1} \left(x ^ {\prime}\right)\right). \tag {3}
$$
This equation is a distributional equivalent to Eq. (1) and thus defines a distributional Bellman operator  $\eta_h^\pi = \mathcal{T}_h^\pi \eta_{h + 1}^\pi$ .
Obviously, from a practical point of view, distributions form a non-parametric family that is not computationally tractable. It is necessary to choose a parametric (thus incomplete) family to represent them. Even the restriction to discrete reward distributions is not tractable, since the number of atoms in the distributions may grow exponentially with the number of steps $^{1}$  [Achab and Neu, 2021]: approximations are unavoidable. The most natural solution is to use projections of the obtained distribution on the parametric family, at each step of the Bellman operator. This process is called parameterization. The practical equivalent to Eq. (1) in DistRL hence writes
$$
\forall x, a, h, \quad \hat {\eta} _ {h} ^ {\pi} (x, a) = \Pi \left(\varrho_ {h} (x, a) * \sum_ {x ^ {\prime}} p _ {h} \left(x, a, x ^ {\prime}\right) \hat {\eta} _ {h + 1} ^ {\pi} \left(x ^ {\prime}, \pi_ {h + 1} \left(x ^ {\prime}\right)\right)\right), \tag {4}
$$
where  $\Pi$  is the projector operator on the parametric family. The full policy evaluation algorithm in DistRL is summarized in Alg.1.
Algorithm 1 Policy Evaluation (Dynamic Programming) for Distributional RL
Input: model  $p$ , reward distributions  $\varrho_h$ , policy  $\pi$  to evaluated,  $\Pi$  projection.
Data:  $\eta \in \mathbb{R}^{H|\mathcal{X}||\mathcal{A}|N}$
$$
\forall x, a \in \mathcal {X} \times \mathcal {A}, \eta_ {H} (x, a) = \delta_ {0}
$$
for  $h = H - 1\to 0$  do
$$
\eta_ {h} (x, a) = \varrho_ {h} (x, a) * \sum_ {x ^ {\prime}} p _ {h} (x, a, x ^ {\prime}) \eta_ {h + 1} \left(x ^ {\prime}, \pi_ {h + 1} \left(x ^ {\prime}\right)\right) \quad \forall x, a \in \mathcal {X} \times \mathcal {A}
$$
$$
\eta_ {h} (x, a) = \Pi \left(\eta_ {h} (x, a)\right) \quad \forall x, a \in \mathcal {X} \times \mathcal {A}
$$
Output:  $\eta_h(x,a)\forall x,a,h$
Distribution Parametrization The most commonly used parametrization is the so-called quantile projection. It put Dirac's (atoms) with fixed weights at locations that correspond to the quantiles of the source distribution. One main benefit is that it does not require a previous knowledge of the support of the distribution, and allows for unbounded distributions.
The quantile projection is defined as
$$
\Pi_ {Q} \mu = \frac {1}{N} \sum_ {i = 0} ^ {N - 1} \delta_ {z _ {i}}, \quad \text {w i t h} (z _ {i}) _ {i} \text {c h o s e n s u c h a s} F _ {\mu} (z _ {i}) = \frac {2 i + 1}{2 N}, \tag {5}
$$
which corresponds to a minimal  $W_{1}$  distance:  $\Pi_{Q}\mu \in \arg \min_{\hat{\mu} = \sum_{i}\delta_{i} / N}W_{1}(\mu ,\hat{\mu})$ , where  $W_{1}(.,.)$  is the Wasserstein distance defined for any distributions  $\nu_{1},\nu_{2}$  as  $W_{1}(\nu_{1},\nu_{2}) = \int_{0}^{1}\left|F_{\nu_{1}}^{-1}(u) - F_{\nu_{2}}^{-1}(u)\right|\mathrm{d}u$ . Note that this parametrization might admit several solutions and thus the projection may not be unique. For simplicity, we overload the notation to  $\Pi_{Q}\eta = (\Pi_{Q}\eta (x,a))_{(x,a)\in \mathcal{X}\times \mathcal{A}}$
1The support of the return (sum of the rewards) is incremented at each step by a number of atoms that depend on the current support.
For a Q-value distribution  $\eta$  with support of length  $\Delta_{\eta}$ , and parametrization of resolution  $N$ , Rowland et al. [2019] prove that the projection error is bounded by
$$
\sup  _ {(x, a) \in \mathcal {X} \times \mathcal {A}} W _ {1} \left(\Pi_ {Q} \eta (x, a), \eta (x, a)\right) \leq \frac {\Delta_ {\eta}}{2 N}. \tag {6}
$$
In Section 3, we extend this result to the full iterative Policy Evaluation process and bound the error on the returned statistical functional in the finite-horizon setting. Note that other studied parametrizations exist but are less practical. For completeness, we discuss the Categorical Projection [Bellemare et al., 2017][Rowland et al., 2018][Bellemare et al., 2023] in Appendix B.
2.2 Beyond expected returns
The expected value is an important functional of a probability distribution, but it is not the only one of interest in decision theory - especially when a control of the risk is important. We discuss two concepts that have received considerable attention: utilities, defined as expected values of functions of the return, and distorted means which place emphasis on certain quantiles.
Expected Utilities are of the form  $\mathbb{E}[f(Z)]$ , or  $\int f \, \mathrm{d}\nu$ , where  $Z$  is the return of distribution  $\nu$  and  $f$  is an increasing function. For instance, when  $f$  is a power function, we obtain the different moments of the return. The case of exponential functions plays a particularly important role: the resulting utility is referred to as exponential utility, exponential risk measure, or generalized mean according to the context:
$$
U _ {\exp} (\nu) = \frac {1}{\lambda} \log \mathbb {E} \left[ \exp (\lambda X) \right] \quad X \sim \nu \text {a n d} \lambda \in \mathbb {R}. \tag {7}
$$
This family of utilities has a variety of applications in finance, economics, and decision making under uncertainty [Föllmer and Schied, 2016]. They can be considered as a risk-aware generalization of the expectation, with benefits such as accommodating a wide range of behaviors [Shen et al., 2014] from risk-seeking when  $\lambda > 0$ , to risk-averse when  $\lambda < 0$  (the limit  $\lambda \to 0$  is exactly the expectation). To fix ideas,  $U_{\mathrm{exp}}\bigl (\mathcal{N}(\mu ,\sigma^2)\bigr) = \mu +\lambda \sigma^2$ : each  $\lambda$  captures a certain quantile of the Gaussian distribution.
Distorted means, on the other hand, involve taking the mean of a random variable, but with a different weighting scheme [Dabney et al., 2018a]. The goal is to place more emphasis on certain quantiles, which can be achieved by considering the quantile function  $F^{-1}$  of the random variable and a continuous increasing function  $\beta : [0,1] \to [0,1]$ . By applying  $\beta$  to a uniform variable  $\tau$  on  $[0,1]$  and evaluating  $F^{-1}$  at the resulting value  $\beta(\tau)$ , we obtain a new random variable that takes the same values as the original variable, but with different probabilities. The distorted mean is then calculated as the mean of this new random variable, given by the formula  $\int \beta'(\tau) F^{-1}(\tau) d\tau$ . If  $\beta$  is the identity function, the result is the classical mean. When  $\beta$  is  $\tau \mapsto \min(\tau / \alpha, 1)$ , we get the  $\alpha$ -Conditional Value at Risk (CVaR(α)) of the return, a risk measure widely used in risk evaluation [Rockafellar et al., 2000].
3 Policy Evaluation
The theory of MDPs is particularly developed for estimating and optimizing the mean of the return of a policy. But other values associated to the return can be computed the same way, by dynamic programming. This includes for instance the variance of the return, or more generally, any moment of order  $p \geq 2$ , as was already noticed in the 1980's [Sobel, 1982]. Recently, Rowland et al. [2019] showed that for utilities in discounted MDPs, this is essentially all that can be done. More precisely, they introduce the notion of Bellman closedness (recalled below for completeness) that characterizes a finite set of statistics that can efficiently be computed by dynamic programming.
Definition 1 (Bellman closedness [Rowland et al., 2019]). A set of statistical functionals  $\{s_1, \ldots, s_K\}$  is said to be Bellman closed if for each  $(x, a) \in \mathcal{X} \times \mathcal{A}$ , the statistics  $s_{1:K}(\eta_h^\pi(x, a))$  can be expressed in closed form in terms of the random variables  $R_h(x, a)$  and  $(s_{1:K}(\eta_{h+1}^\pi(X', A'))$ ,  $A' \sim \pi(x)$ ,  $X' \sim p_h(x, A', \cdot)$ , independently of the MDP.
Importantly, in the undiscounted setting, Rowland et al. [2019](Appendix B, Theorem 4.3) show that the only families of utilities that are Bellman closed are of the form  $\{x\mapsto x^{\ell}\exp (\lambda x)|0\leq \ell \leq L\}$
for some  $L < \infty$ . Thus, all utilities and statistics of the form of (or linear combinations of) moments and exponential utilities can easily be computed by classic linear dynamic programming and do not require distributional RL (see Appendix A.3).
Some important metrics such as the CVaR or the quantiles are not known to belong to any Bellman-closed set and hence cannot be easily computed. For this kind of function of the return, the knowledge of the transitions and the values in following steps is insufficient to compute the value on a specific step. In general, it requires the knowledge of the whole distribution of each reward in each state. Hence, techniques developed in distributional RL come in handy: for a choice of parametrization, one can use the projected dynamic programming step Eq. (4) to propagate a finite set of values along the MDP and approximate the distribution of the return. In the episodic setting, following the line of Rowland et al. [2019] (see Eq.(6)), we prove that the Wasserstein distance error between the exact and approximate distribution of the Q-values of a policy is bounded.
Proposition 1. Let  $\pi$  be a policy and  $\eta^{\pi}$  the associated  $Q$ -value distributions. Assume the return is bounded on a interval of length  $\Delta_{\eta} \leq H\Delta_{R}$ , where  $\Delta_{R}$  is the support size of the reward distribution. Let  $\hat{\eta}^{\pi}$  be the  $Q$ -value distributions obtained by dynamic programming (Algorithm 1) using the quantile projection  $\Pi_{Q}$  with resolution  $N$ . Then,
$$
\sup _ {(x, a, h) \in (\mathcal {X}, \mathcal {A}, [ H ])} W _ {1} (\hat {\eta} _ {h} ^ {\pi} (x, a), \eta_ {h} ^ {\pi} (x, a)) \leq H \frac {\Delta_ {\eta}}{2 N} \leq H ^ {2} \frac {\Delta_ {R}}{2 N}.
$$
This result shows that the loss of information due to the parametrization may only grow quadratically with the horizon. The proof consists of summing the projection bound in (6) at each projection step, and using the non-expansion property of the Bellman operator [Bellemare et al., 2017]. The details can be found in Appendix C
The key question is then to understand how such error translates into our estimation problem when we apply the function of interest to the approximate distribution. We provide a first bound on this error for the family of statistics that are either utilities or distorted means.
First, we prove that the utility is Lipschitz on the set of return distributions.
Lemma 1. Let  $s$  be either an utility or a distorted mean and let  $L$  be the Lipschitz coefficient of its characteristic function. Let  $\nu_{1}, \nu_{2}$  be return distributions. Then:
$$
\left| s \left(\nu_ {1}\right) - s \left(\nu_ {2}\right) \right| \leq L W _ {1} \left(\nu_ {1}, \nu_ {2}\right).
$$
Both family of functionals are treated separately, but lays a similar bound. The utility bound is the direct application of the Kantorovitch-Rubenstein duality, while the distorted mean one is a direct majoration in the integral. Again, the details are provided in the Appendix.
This property allows us to prove a maximal upper bound on the estimation error for those two families.
Theorem 1. Let  $\pi$  be a policy. Let  $\eta^{\pi}$  be the  $Q$ -value return distribution associated to  $\pi$  with the return bounded on a interval of length  $\Delta_{\eta} \leq H\Delta_{R}$  where  $\Delta_{R}$  is the support size of the reward distribution. Let  $\hat{\eta}^{\pi}$  be the approximated return distribution computed with Algorithm 1, for the projection  $\Pi_{Q}$  with resolution  $N$ . Let  $s$  be either an expected utility or a distorted mean, and  $L$  the Lipschitz coefficient of its characteristic function. Then:
$$
\sup  _ {x, a, h} | s (\hat {\eta} _ {h} ^ {\pi} (x, a)) - s (\hat {\eta} _ {h} ^ {\pi} (x, a)) | \leq L H \frac {\Delta_ {\eta}}{2 N} \leq L H ^ {2} \frac {\Delta_ {R}}{2 N}.
$$
Note that depending on the choice of utilities, the Lipschitz coefficient  $L$  may also depend on  $H$  and  $\Delta R$ . For instance, in a stationary MDP, the Lipschitz constant of the exponential utility depends exponentially on  $\Delta_{\eta}$ . For the CVaR(α), however,  $L$  is constant and only depends on  $\alpha \in (0,1)$ .
Experiment: empirical validation of the bounds on a simple MDP We consider a simple Chain MDP environment of length  $H = 70$  equal to the horizon (see Figure 1 (right)) [Rowland et al., 2019], with a single action leading to the same discrete reward distribution for every step. We consider a Bernoulli reward distribution  $\mathcal{B}(0.5)$  for each state so that the number of atoms for the return only grows linearly $^2$  with the number of steps, which allows to compute the exact distribution easily.
2At round  $h\in [H]$  , the support of the return is  $\{0,1,\dots,h\}$  , so  $h$  atoms.
Figure 1: A Chain MDP of length  $H$  with deterministic transition and identical reward distribution for each state.
Figure 2: Left: Validation of Theorem 1 on  $\mathrm{CVaR}(\alpha)$  together with the scaled upper bound (see main text for discussion): the quadratic dependence in  $H$  is verified. Right: Validation of Proposition 1: The cumulative projection error (dashed blue) is the sum of the projection errors at every time step, and matches the true approximation error (solid blue). The theoretical upper bound (dashed red) differs only by a factor 2.
We compare the distributions obtained with exact dynamic programming and the approximate distribution obtained by Alg 1, with a quantile projection with resolution  $N = 1000$ . Note that even at early stages, when the true distribution has less atoms than the resolution, the exact and approximate distributions differ due to the weights of the atoms in the quantile projection. Figure 2 (Right) reports the Wasserstein distance between the two distributions: the cumulative projection approximation error (dashed blue), the true error between the current exact and approximate distributions (solid blue) and the theoretical bound (red). Fundamentally, the proof of Prop. 1 upper bounds the distance between distributions by the cumulative projection error so we plot this quantity to help validating it.
We also empirically validate Theorem 1 by computing the  $\mathrm{CVaR}(\alpha)$  for  $\alpha \in \{0.1, 0.25\}$ , corresponding respectively to distorted means with Lipschitz constants  $L = \{10, 4\}$ . We compute these statistics for both distributions and report the maximal error together with the theoretical bound, re-scaled by a factor 5. Figure 2 (Left) shows an impressive correspondence of the theory and the empirical results despite a constant multiplicative gap.
4 Planning
Planning refers to the problem of returning a policy that optimizes our objective for a given model and reward function (or distribution in DistRL). It shares with policy evaluation the property to be grounded on a Bellman equation: see Eq. (2) for the classical expected return, which leads to efficient computations by dynamic programming.
For other statistical functionals of the cumulated reward, however, can the optimal policy be computed efficiently? The main result of this section characterizes the family of functionals that can be exactly and efficiently optimized by Dynamic Programming.
In the previous section, we recalled that Bellman-closed sets of utilities can be efficiently computed by DP as long as all the values of the utilities in the Bellman-closed family are computed together. For the planning problem, however, we only want to optimize one utility so we cannot consider families as previously. Under this constraint, only exponential and linear expected utilities are Bellman closed and thus can verify a Bellman Equation. In fact, for the exponential utilities, such Bellman Equation