Variance Control for Distributional Reinforcement Learning

Consider a finite Markov Decision Process (MDP) $(\mathcal{X},\mathcal{A},P,\gamma,\mathcal{R})$, with a finite set of states $\mathcal{X}$, a finite set of actions $\mathcal{A}$, the transition kernel $P:\mathcal{X}\times\mathcal{A}\rightarrow\mathscr{P}(\mathcal{X})$, the discounted factor $\gamma\in[0,1)$, and the bounded reward function $\mathcal{R}:\mathcal{X}\times\mathcal{A}\rightarrow\mathscr{P}([-R_{max},R_{max}])$. At each timestep, an agent observes state $X_{t}\in\mathcal{X}$, takes an action $A_{t}\in\mathcal{A}$, transfers to the next state $X_{t+1}\sim P\left(\cdot\mid X_{t},A_{t}\right)$, and receives a reward $R_{t}\sim\mathcal{R}\left(X_{t},A_{t}\right)$. The state-action value function $Q^{\pi}:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}$ of a policy $\pi:\mathcal{X}\rightarrow\mathscr{P}(\mathcal{A})$ is the expected discounted sum of rewards starting from $x$, taking an action $a$ and following a policy $\pi$. $\mathscr{P}(\mathcal{X})$ denotes the set of probability distributions on a space $\mathcal{X}$.

The classic Bellman equation (Bellman, 1966) relates expected return at each state-action pair $(x,a)$ to the expected returns at possible next states by:

In the learning task, Q-Learning (Watkins, 1989) employs a common way to obtain $\pi^{*}$, which is to find the unique fixed point $Q^{*}=Q^{\pi^{*}}$ of the Bellman optimality equation:

Instead of directly estimating the expectation $Q^{\pi}(x,a)$, DRL focuses on estimating the distribution of the sum of discounted rewards $\eta_{\pi}(x,a)=\mathcal{D}(\sum_{t=0}^{\infty}\gamma^{t}R_{t}\mid X_{0}=x,A_{0}=a)$ to sufficiently capture the intrinsic randomness, where $\mathcal{D}$ extract the probability distribution of a random variable. In analogy with Equation 2, $\eta_{\pi}$ satisfies the distributional Bellman equation (Bellemare et al., 2017) as follows,

where $f_{\gamma,r}:\mathbb{R}\to\mathbb{R}$ is defined by $f_{\gamma,r}(x)=r+\gamma x,$ and $\left(f_{\gamma,r}\right)_{\#}\eta$ is the pushforward measure of $\eta$ by $f_{\gamma,r}$. Note that $\eta_{\pi}$ is the fixed point of distributional Bellman operator $\mathcal{T}^{\pi}:\mathscr{P}(\mathbb{R})^{\mathcal{X}\times\mathcal{A}}\to\mathscr{P}(\mathbb{R})^{\mathcal{X}\times\mathcal{A}}$, i.e., $\mathcal{T}^{\pi}\eta_{\pi}=\eta_{\pi}$.

In general, the return distribution supports a wide range of possible returns and its shape can be quite complex. Moreover, the transition dynamics are usually unknown in practice, and thus the full computation of the distributional Bellman operator is usually either impossible or computationally infeasible. In the following subsections, we review two main categories of DRL algorithms relying on parametric approximations and projection operators.

We first show the convergence of both the expectation and the variance of the distributional Bellman operator $\mathcal{T}^{\pi}$. Then, we take parametric representation and projection operator into consideration.

Based on the fact that $\mathcal{T}^{\pi}$ is a $\gamma$-contraction in $\bar{d}_{p}$ metric (Bellemare et al., 2017), where $\bar{d}_{p}$ is the maximal form of the Wasserstein metric, Proposition 3.1 implies that $\mathcal{T}^{\pi}$ is a contraction for both the expectation and the variance. The two converge exponentially to their true values by iteratively applying the distributional Bellman operator.

However, in practice, employing parametric representation for the return distribution leaves a theory-practice gap, which makes neither the expectation nor the variance converge to the true values. To better understand the bias in the Q-function approximation caused by the parametric representation, we introduce the concept of mean-preserving ¹¹1This property has been thoroughly discussed in previous work. Based on Section 5.11 of Bellemare et al. (2023), we conclude this definition. to describe the relationship between the expectations of the original distribution and the projected distribution:

For CDRL, a discussion of the mean-preserving property is given by Lyle et al. (2019) and Rowland et al. (2019). It can be shown that for any $\nu\in\mathscr{F}_{\mathcal{C}}$, where $\mathscr{F}_{\mathcal{C}}$ is a $N$-categorical representation, the projection $\Pi_{\mathcal{C}}$ preserves the distribution’s expectation when its support is contained in the interval $[z_{1},z_{N}]$. However, these practitioners usually employ a wide predefined interval for return which makes the projection operator typically overestimate the variance.

For QDRL, $\Pi_{\mathcal{W}_{1}}$ is not mean-preserving (Bellemare et al., 2023). Given any distribution $\nu\in\mathscr{F}_{\mathcal{W}_{1}}$, where $\mathscr{F}_{\mathcal{W}_{1}}$ is a $N$-quantile representation, there is no unique $N$-quantile distribution $\Pi_{\mathcal{W}_{1}}\nu$ in most cases, as the projection operator $\Pi_{\mathcal{W}_{1}}$ is not a non-expansion in 1-Wasserstein distance (See Appendix B for details). This means that the expectation, variance, and higher-order moments are not preserved. To make this concrete, a simple MDP example is used to illustrate the bias in the learned quantile estimates.

In Figure 2 (a), rewards $R_{1}$ and $R_{2}$ are randomly sampled from $\mathrm{Unif}(0,1)$ and $\mathrm{Unif}(1/N,1+1/N)$ at states $x_{1}$ and $x_{2}$ respectively, and no rewards are received at $x_{0}$. Clearly, the true return distribution at state $x_{0}$ is the mixture $\frac{\gamma}{2}(R_{1}+R_{2})$, hence the $\frac{1}{2N}$-th quantile is $\frac{\gamma}{N}$. When using the QDRL algorithm with $N$ quantile estimates, the approximated return distribution $\hat{\eta}(x_{1},a)=\frac{1}{N}\sum_{i=1}^{N}\delta_{\frac{2i-1}{2N}}$ and $\hat{\eta}(x_{2},a)=\frac{1}{N}\sum_{i=1}^{N}\delta_{\frac{2i+1}{2N}}$. In this case, the $\frac{1}{2N}$-th quantile of the approximated return distribution at state $x_{0}$ is $\frac{3\gamma}{2N}$, whereas the true value is $\frac{\gamma}{N}$. Moreover, for each $i=1,\ldots,N$, the $\frac{2i-1}{2N}$-th quantile estimate at state $x_{0}$ is not equal to the true value.

These biased quantile estimates illustrated in Figure 2 (a) are caused by the use of quantile representation and its projection operator $\Pi_{\mathcal{W}_{1}}$. This undesirable property in turn affects the QDRL update, as the combined operator $\Pi_{\mathcal{W}_{1}}\mathcal{T}^{\pi}$ is in general not a non-expansion in $\bar{d}_{p}$, for $p\in[1,\infty)$ (Dabney et al., 2018b), which means that the learned quantile estimates may not converge to the true quantiles of the return distribution ²²2 A recent study (Rowland et al., 2023a) proves that QDRL update may have multiple fixed points, indicating quantiles may not converge to the truth. Despite this, Proposition 2 (Dabney et al., 2018b) concludes that the projected Bellman operator $\Pi_{\mathcal{W}_{1}}\mathcal{T}^{\pi}$ remains a contraction in $\bar{d}_{\infty}$. This implies that quantile convergence is guaranteed for all $p\in[1,\infty]$.. The projection operator $\Pi_{\mathcal{W}_{1}}$ is not mean-preserving which inevitably leads to bias in the expectation of return distribution when iteratively applying the projected Bellman operator $\Pi_{\mathcal{W}_{1}}\mathcal{T}^{\pi}$ during the training process, resulting in a deviation between the estimate and the true value of the Q-function in the end. We now derive an upper bound to quantify this deviation, i.e. $\mathcal{E}_{3}$.

Theorem 3.3 implies that the convergence of expectation with projected Bellman operator $\Pi_{\mathcal{W}_{1}}\mathcal{T}^{\pi}$ cannot be guaranteed after quantile representation and its projection operator are applied ³³3Note that this bound has a limitation, which only considers the one-step effect of applying the projection operator $\Pi_{W_{1}}$. Therefore, it becomes irrelevant with the iteration number $k$. However, Proposition 4.1 of Rowland et al. (2023b) provides a more compelling bound considering the cumulative effect of iteratively applying $\Pi_{W_{1}}$.. Note that the bound will tend to zero with $N\rightarrow\infty$, thus it is reasonable to use a relatively large representation size $N$ to reduct $\mathcal{E}_{3}$ in practice.

The other two types of errors $\mathcal{E}_{1}$ and $\mathcal{E}_{2}$, which determine the variance of the Q-function estimate, are accumulated during the training process by keeping encountering unseen state-action pairs. The target approximation error $\mathcal{E}_{1}$ affects action selections, while the Bellman operator approximation error $\mathcal{E}_{2}$ leads to the accumulated error of the Q-function estimate, which can be amplified by using the temporal difference updates (Sutton, 1988). The accumulated errors of the Q-function estimate with high uncertainty can make some certain states to be incorrectly estimated, leading to suboptimal policies and potentially divergent behaviors.

As depicted in Figure 2 (b), we utilize this toy example to illustrate how QDRL fails to learn an optimal policy due to a high variance of the approximation error. This 5-state MDP example is originally introduced in Figure 7 of Rowland et al. (2019). In this case, $\eta(x_{0},a_{1})$ and $\eta(x_{0},a_{2})$ follow exponential distributions, and the expectations of them are 1.2 and 1, respectively. We consider a tabular setting, which uniquely represents the approximated return distribution at each state-action pair. Figure 2 (c) demonstrates that in policy evaluation, QDRL inaccurately approximates the Q-function, as it underestimates the expectation of $\eta(x_{0},a_{1})$ and overestimates the other. This is caused by the poor capture of tail events, which results in high uncertainty in the Q-function estimate. Due to the high variance, QDRL fails to learn the optimal policy and chooses a non-optimal action $a_{2}$ at the initial state $x_{0}$. On the contrary, our proposed algorithm, QEMRL, employs a statistically robust estimator of the Q-function to reduce its variance, relieves the underestimation and overestimation issues, and ultimately allows for more efficient policy learning.

Different from previous QDRL studies that focus on exploiting the distribution information to further improve the model performance, this work highlights the importance of controlling the variance of the approximation error to obtain a more accurate estimate of the Q-function. More discussion about this is given in the following section.

In the context of quantile-based DRL, Q-function is the integral of the quantiles. To approximate this, QDRL employs a simple empirical mean (EM) estimator $\frac{1}{N}\Sigma_{i}\hat{q}(\tau_{i})$, and it is natural to assume that the estimated quantile satisfies

where $\varepsilon(\tau)$ is a zero-mean error. In this case, considering the crossing issue and the biased tail estimates, we assume that the variance of $\varepsilon(\tau)$ is non-constant and depends on $\tau$, which is usually called heteroskedasticity in statistics.

For a direct understanding, we conduct a simple simulation using a Chain MDP to illustrate how QDRL can fail to fit the quantile function. As shown in Figure 3(b), QDRL fits well in the peak area but struggles at the bottom and the tail. Moreover, the non-monotonicity of the quantile estimates in the poorly fitted areas is more severe than the others. As the deviations of the quantile estimates from the truths is significantly larger in the low probability region and the tail, we can make the heteroskedasticity assumption in this case. This phenomenon can be explained since samples near the bottom and the tail are less likely to be drawn. In real-world situations, multimodal distributions are commonly encountered and the heteroskedasticity problem may result in imprecise distribution approximations and consequently poor Q-function approximations. In the next part, we will discuss how to enhance the stability of the Q-function estimate.

It is well-known that quantile can be expressed by the Cornish-Fisher Expansion (CFE, Cornish & Fisher, 1938):

where $z_{\tau}$ is the $\tau$-th quantile of the standard normal distribution, $\mu$ is the mean, $\sigma$ is the standard deviation, $s$ and $k$ are the skewness and kurtosis of the interested distribution, and the remaining terms in the ellipsis are higher-order moments (See Appendix C for more details). The CFE theoretically determines the distribution with known moments and is widely used in financial studies. Recently, Zhang & Zhu (2023) employ CFE to estimate higher-order moments of financial time series data, which are not directly observable. Our method utilizes a truncated version of CFE framework and employs a linear regression model to construct efficient estimators for distribution moments based on known quantiles. Consequently, we apply this approach within the context of quantile-based DRL.

To be more specific, we plug in the estimate $\hat{q}({\tau})$ of the the $\tau$-th quantile to Equation 5 and expand it by the first order:

where $m_{1}$ is the mean (say, $1$-th moment) of the return distribution, i.e., the Q-function, and $\omega_{1}(\tau)$ is the remaining term associated with the higher-order ($>1$-th) moments. If $\omega_{1}(\tau)$ is negligible, $m_{1}$ can be estimated by averaging the $N$ quantile estimates in QDRL.

When the estimated quantile is expanded to the second order, we particularly have the following representation:

where $\omega_{2}(\tau)$ is the remaining term associated with the higher-order ($>2$-th) moments. Assume that $\omega_{2}(\tau)$ is negligible, we can derive a regression model by plugging in the $N$ quantile estimates, such that

The higher-order expansions can be conducted in the same manner. Note that the remaining term is omitted for constructing a regression model, and a more in-depth analysis of the remaining term is available in Section C.2.

For notation simplicity, we rewrite (22) in a matrix form,

where $\boldsymbol{\hat{Q}}\in\mathbb{R}^{N}$ is the vector of estimated quantiles, $\mathbf{X}_{2}\in\mathbb{R}^{N\times 2}$ and $\boldsymbol{M}_{2}\in\mathbb{R}^{2}$ are the design matrix and the moments respectively, and $\mathcal{E}$ is the vector of error terms.

For this bivariate regression model (23), the traditional ordinary least squares method (OLS) can be used to estimate $\boldsymbol{M}_{2}=(m_{1},\sqrt{m_{2}})^{\prime}$ when the variances of the errors are invariant across different quantile locations, also known as the homoscedasticity assumption. The estimator $\hat{m}_{1}$ is denoted as Quantiled Expansion Mean (QEM) in this work. However, since the homoscedasticity assumption required by OLS is always violated in real cases, we may consider using the weighted ordinary least squares method (WLS) instead. Under the normality assumption, the following results tell that the WLS estimator $\hat{m}_{1}$ has a lower variance than the direct empirical mean.

Remark: Note that it is impossible to determine the weight matrix $V$ for each state-action pair in practice. Hence, we focus on capturing the relatively high variance in the tail, specifically in the range of $\tau\in(0,0.1]\cup[0.9,1)$. To achieve this, we use a constant $v_{i}$, which is set to a value greater than 1 in the tail and equal to 1 in the rest. $v_{i}$ is treated as a hyperparameter to be tuned in practice (See Appendix E).

With Lemma 4.1, the reduction of variance can be guaranteed by the following Proposition 4.2. Throughout the training process, heteroskedasticity is inevitable, and thus the QEM estimator always exhibits a lower variance than the standard EM estimator $\hat{m}^{*}_{1}=\frac{1}{N}\sum_{i=1}^{N}\hat{q}(\tau_{i})$.

We also try to explore the potential benefits of the variance reduction technique QEM in improving the approximation accuracy. The Q-function estimate with higher variance can lead to noisy policy gradients in policy-based algorithms (Fujimoto et al., 2018) and prevent selection optimal actions in value-based algorithms (Anschel et al., 2017). These issues can slow down the learning process and negatively impact the algorithm performance. By the following theorem, we are able to show that QEM can reduce the variance and thus improve the approximation performance.

Theorem 4.3 indicates that a lower concentration bound can be obtained with a smaller $\alpha$ value. The decrease in $\alpha$ can be attributed to the benefits of QEM. Specifically, QEM helps to decrease the perturbations on the Q-function and reduce the variance of the policy gradients, which allows for faster convergence of the policy training and a more accurate distribution approximation. To conclude, QEM relieves the error accumulation within the Q-function update, improves the estimation accuracy, reduces the risk of underestimation and overestimation, and thus ultimately enhances the stability of the whole training process.

FrozenLake (Brockman et al., 2016) is a classic benchmark problem for Q-learning control with high stochasticity and sparse rewards, in which an agent controls the movement of a character in an $n\times n$ grid world. As shown in Figure 4 with a FrozenLake-$4\times 4$ task, ”S” is the starting point, ”H” is the hole that terminates the game, ”G” is the goal state with a reward of 1. All the blue grids stand for the frozen surface where the agent can slide to adjacent grids based on some underlying unknown probabilities when taking a certain movement direction. The reward received by the agent is always zero unless the goal state is reached.

We first approximate the return distribution under the optimal policy $\pi^{*}$, which can be realized using the value iteration approach. To be specific, we start from the ”S” state and perform 1K Monte-Carlo (MC) rollouts. An empirical distribution can be obtained by summarizing all these recording trajectories. With the approximation of the distribution, we can draw a curve of quantile estimates shown in Figure 5. Both QEMRL and QDRL were run for 150K training steps and the $\epsilon$-greedy exploration strategy is applied in the first 1K steps. For both methods, we set the total number of quantiles to be $N=128$.

Although both QEMRL and QDRL can eventually find the optimal movement at the start state, their approximations of the return distribution are quite different. Figure 4 (b) visualizes the approximation errors of the Q-function and the distribution for QEMRL and QDRL with respect to the number of training steps. The Q-function estimates of QEMRL converge correctly in average, whereas the estimates of QDRL do not converge exactly to the truth. A similar pattern can also be found when it comes to the distribution approximation error. Besides, the reduction of variance by using QEM can be verified by the fact that the curves of QEMRL are more stable and decline faster. In Figure 4 (c), we show that the distribution at the start state estimated by QEMRL is eventually closer to the ground truth.

We do some experiments using the MuJoCo benchmark to further verify the analysis results in Section 4. Our implementation is based on the Distributional Soft Actor-Critic (DSAC, Ma et al., 2020) algorithm, which is a distributional version of SAC. Figure 5 demonstrate that both DSAC and QEM-DSAC significantly outperform the baseline SAC. Among the two, QEM-DSAC performs better than DSAC and the learning curves are more stable, which demonstrates that QEM-DSAC can achieve a higher sample efficiency.

We also do some comparison between QEM and the baseline method QR-DQN on the Atari 2600 platform. Figure 8 plots the final results of these two algorithms in six Atari games. At the early training stage, QEM-DQN exhibits significant gain in sampling efficiency, resulting in faster convergence and better performance.

Extension to IQN. Some great efforts have been made by the community of DRL to more precisely parameterize the entire distribution with a limited number of quantile locations. One notable example is the introduction of Implicit Quantile Networks (IQN, Dabney et al., 2018a), which tries to recover the continuous map of the entire quantile curve by sampling a different set of quantile values from a uniform distribution $\mathrm{Unif}(0,1)$ each time.

Our method can also be applied to IQN as it uses the EM approach to estimate the Q-function. It is noted that the design matrix $\mathbf{X}$ must be updated after re-sampling all the quantile fractions at each training step. Moreover, one important sufficient condition $z_{\tau_{i}}=-z_{\tau_{N-i}}$ which ensures the reduction of variance does not hold in the IQN case as $\tau$’s are sampled from a uniform distribution. However, according to the simulation results in LABEL:table4, the variance reduction still remains valid in practice. In this case, all the baseline methods are modified to the IQN version. As Figure 6 and Figure 7 demonstrate, QEM can achieve some performance gain in most scenarios and the convergence speeds can be slightly increased.

Since QEM also provides an estimate of the variance, we may consider using it to develop an efficient exploration strategy. In some recent study studies, to more sufficiently utilize the distribution information, Mavrin et al. (2019) proposes a novel exploration strategy, Decaying Left Truncated Variance (DLTV) by using the left truncated variance of the estimated distribution as a bonus term to encourage exploration in unknown states. The optimal action $a^{*}$ at state $x$ is selected according to $a^{*}=\arg\max_{a^{\prime}}\left(Q\left(x,a^{\prime}\right)+c_{t}\sqrt{\sigma^{2}_{+}}\right)$, where $c_{t}$ is a decay factor to suppress the intrinsic uncertainty, and $\sigma^{2}_{+}$ denotes the estimation of variance. Although DLTV is effective, the validity of the computed truncation lacks a theoretical guarantee. In this work, we follow the idea of DLTV and examine the model performance by using either the variance estimate obtained by QEM or the original DLTV estimation in some hard-explored games. As Figure 8 shows, by using QEM, the exploration efficiency is significantly improved compared to QR-DQN+DLTV since QEM enhances the accuracy of the quantile estimates and thus the accuracy of the distribution variance.