The Nature of Temporal Difference Errors in
Multi-step Distributional Reinforcement Learning

In general, the return $G^{\pi}(x,a)$ is a random variable and we define its distribution as $\eta^{\pi}(x,a)\coloneqq\text{Law}_{\pi}\left(G^{\pi}(x,a)\right)$. The return distribution satisfies the distributional Bellman equation (morimura2010nonparametric, ; morimura2012parametric, ; bellemare2017distributional, ; rowland2018analysis, ; bdr2022, ),

where $(\textrm{b}_{r,\gamma})_{\#}:\mathscr{P}(\mathbb{R})\rightarrow\mathscr{P}(\mathbb{R})$ is the pushforward operation defined through the function $\textrm{b}_{r,\gamma}(z)=r+\gamma z$ (rowland2018analysis, ). For convenience, we adopt the notation $\eta^{\pi}(X_{t},A_{t}^{\pi})\coloneqq\sum_{a}\pi(a|X_{t})\eta^{\pi}(X_{t},a)$. Throughout the paper, we focus on the space of distributions with bounded support $\mathscr{P}_{\infty}(\mathbb{R})$. Let $\eta\in\mathscr{P}_{\infty}(\mathbb{R})^{\mathcal{X}\times\mathcal{A}}$ be any distribution vector, we define the distributional Bellman operator $\mathcal{T}^{\pi}:\mathscr{P}_{\infty}(\mathbb{R})^{\mathcal{X}\times\mathcal{A}}\rightarrow\mathscr{P}_{\infty}(\mathbb{R})^{\mathcal{X}\times\mathcal{A}}$ as follows (rowland2018analysis, ; bdr2022, ),

Let $\eta^{\pi}$ be the collection of return distributions under $\pi$; the distributional Bellman equation can then be rewritten as $\eta^{\pi}=\mathcal{T}^{\pi}\eta^{\pi}$. The distributional Bellman operator $\mathcal{T}^{\pi}$ is $\gamma$-contractive under the supremum $p$-Wasserstein distance (dabney2018distributional, ; bdr2022, ), so that $\eta^{\pi}$ is the unique fixed point of $\mathcal{T}^{\pi}$. See Appendix B for details of the distance metrics.

We provide a brief background on the value-based multi-step off-policy setting as a reference for the distributional case discussed below. In off-policy learning, the data is generated under a behavior policy $\mu$, which potentially differs from target policy $\pi$. The aim is to evaluate the target Q-function $Q^{\pi}$. As a standard assumption, we require $\text{supp}(\pi(\cdot|x))\subseteq\text{supp}(\mu(\cdot|x)),\forall x\in\mathcal{X}$. Let $\rho_{t}\coloneqq\pi(A_{t}|X_{t})/\mu(A_{t}|X_{t})$ be the step-wise importance sampling (IS) ratio at time step $t$. Step-wise IS ratios are critical in correcting for the off-policy discrepancy between $\pi$ and $\mu$.

Let $c_{t}\in[0,\rho_{t}]$ be a time-dependent trace coefficient. We denote $c_{1:t}=c_{1}\cdots c_{t}$ and define $c_{1:0}=1$ by convention. Consider a generic form of the return-based off-policy operator $R^{\pi,\mu}$ as in (munos2016safe, ),

In the above and below, we omit the notation conditioning on $X_{0}=x,A_{0}=a$ for conciseness. The general form of $R^{\pi,\mu}$ encompasses many important special cases: when on-policy and $c_{t}=\lambda$, it recovers the Q-function variant of TD($\lambda$) (sutton1998, ; harutyunyan2016q, ); when $c_{t}=\lambda\min(\overline{c},\rho_{t})$, it recovers a specific form of Retrace (munos2016safe, ); when $c_{t}=\rho_{t}$, it recovers the importance sampling (IS) operator. The back-up target is computed as a mixture over TD errors $\delta_{t}^{\pi}$, each calculated from the one-step transition data. We also define the discounted TD error $\widetilde{\delta}_{t}^{\pi}=\gamma^{t}\delta_{t}^{\pi}$, which can be interpreted as the difference between $n$-step returns from two time steps $t$ and $t+1$, as we discuss in Section 4. As we will detail, the property of $\widetilde{\delta}_{t}^{\pi}$ marks a significant difference from the distributional RL setting.

By design, $R^{\pi,\mu}$ has $Q^{\pi}$ as the unique fixed point. Multi-step updates make use of rewards from multiple time steps, propagating learning signals more efficiently. This is reflected by the fact that $R^{\pi,\mu}$ is $\beta$-contractive with $\beta\in[0,\gamma]$ (munos2016safe, ) and often contracts to $Q^{\pi}$ faster than the one-step Bellman operator $T^{\pi}$. Our goal is to design distributional equivalents of multi-step off-policy operators, which can lead to concrete algorithms with sample-based learning.

The value-based Retrace back-up target can be written as a mixture of value-based TD errors. To better parse the distributional Retrace operator and draw comparison to the value-based setting, we seek to rewrite the distributional back-up target $\mathcal{R}^{\pi,\mu}\eta(x,a)$ into a weighted sum of some notion of distributional TD errors. To this end, we start with a natural analogy to the value-based TD error.

When the context is clear, we also adopt the concise notation $\Delta_{t}^{\pi}=\Delta^{\pi}(X_{t},A_{t},R_{t},X_{t+1})$. By construction, distributional TD errors are signed measures with zero total mass (bdr2022, ). The distributional TD error is a natural counterpart to the value-based TD error, because they both stem directly from the corresponding one-step Bellman operators. However, unlike in value-based RL, where TD errors alone suffice to specify the multi-step learning operator (Equation (3)), in distributional RL this is not enough. We introduce the path-dependent distributional TD error, which serves as the building block to distributional Retrace.

Path-dependent distributional TD errors are defined as a pushforward measures from $\Delta_{t}^{\pi}$, where the pushforward operations depend on $G_{0:t-1}$. This equips $\widetilde{\Delta}_{0:t}^{\pi}$ with an intriguing property, path-dependency. Concretely, this means that the path-dependent distributional TD error depends on the sequence of rewards $(R_{s})_{s=0}^{t-1}$ leading up to step $t$. With the above definitions, we can finally rewrite the back-up target of distributional Retrace as a weighted sum of path-dependent distributional TD errors $\mathcal{R}^{\pi,\mu}\eta(x,a)=\eta(x,a)+\mathbb{E}_{\mu}[\sum_{t=0}^{\infty}c_{1:t}\widetilde{\Delta}_{0:t}^{\pi}]$. We now illustrate the difference between value-based and distributional TD errors.

To highlight the difference between distributional and value-based multi-step learning, we discuss the impact that path-dependent distributional TD errors have on the backward-view distributional algorithm. Thus far, distributional back-up targets are expressed in the forward-view, i.e., the back-up target at time $t$ is calculated as a function of future transition tuples $(X_{s},A_{s},R_{s})_{s\leq t}$. The forward-view algorithms, unless truncated, wait until the episode finishes to carry out the update, which might be undesirable when the problem is non-episodic or has a very long horizon.

In the backward-view, when encountering a distributional TD error $\Delta_{t}^{\pi}$, the algorithm carries out updates for all predictions at time $t^{\prime}\leq t$ (sutton1998, ). To this end, the algorithm needs to maintain additional partial return traces, i.e., the partial sum of rewards $G_{t^{\prime}:t}$, in order to calculate the path-dependent TD error $\widetilde{\Delta}_{t}^{\pi}$. Unlike the value-based state-dependent eligibility traces (sutton1998, ; van2020expected, ), partial return traces are time-dependent. This implies that in an episode of $T$ steps, value-based backward-view algorithms require memory of size $\min(|\mathcal{X}||\mathcal{A}|,\mathcal{O}(T))$ while the distributional algorithms requires $\mathcal{O}(T)$.

In addition to the added memory complexity, the incremental updates of distributional algorithms are also much more complicated due to the path-dependent TD errors. We remark that the path-independent nature of value-based TD errors greatly simplify the value-based backward-view algorithm. For a more detailed discussion, see Appendix D.

In our initial derivation, we arrived at $\mathcal{R}^{\pi,\mu}$ through the application of importance sampling (IS) in a different way from the value-based setting. We now highlight the subtle differences and caveats.

For a fixed $n\geq 1$, consider the trace coefficient $c_{t}=\rho_{t}\mathbb{I}[t<n]$. The back-up target of the resulting Retrace operator reduces to $\mathbb{E}_{\mu}\left[\rho_{1:n-1}\cdot\left(\textrm{b}_{G_{0:n-1},\gamma^{n}}\right)_{\#}\eta\left(X_{n},A_{n}^{\pi}\right)\right]$. This can be seen as applying IS to the $n$-step prediction $\left(\textrm{b}_{G_{0:n-1},\gamma^{n}}\right)_{\#}\eta\left(X_{n},A_{n}^{\pi}\right)$. As a caveat, note that an appealing alternative approach is to apply IS to $G_{0:n-1}$, producing the estimate $\left(\textrm{b}_{\rho_{1:n-1}G_{0:n-1},\gamma^{n}}\right)_{\#}\eta\left(X_{n},A_{n}^{\pi}\right)$. This latter estimate does not properly correct for the off-policy discrepancy between $\pi$ and $\mu$. To see why, note that applying the IS ratio to $G_{0:n-1}$, instead of to the probability of its occurrence, is an artifact of value-based RL because the expected return is linear in $G_{0:t}$ (precup2001off, ). In general for distributional RL, one should importance weigh the measures instead of sum of rewards.

Given an input distribution vector $\eta$, we use the distributional Retrace operator to construct the back-up target $\mathcal{R}^{\pi,\mu}\eta$. Then, we use the quantile projection to map the back-up target onto the space of representations $\Pi_{\mathcal{Q}}\mathcal{R}^{\pi,\mu}\eta$. Overall, we are interested in the recursive update: start with any $\eta_{0}\in\mathscr{P}_{\mathcal{Q}}(\mathbb{R})^{\mathcal{X}\times{\mathcal{A}}}$, consider the sequence of distributions generated via $\eta_{k+1}=\Pi_{\mathcal{Q}}\mathcal{R}^{\pi,\mu}\eta_{k}$. A direct application of Proposition 3.2 allows us to characterize the convergence of the sequence, following the approach of bdr2022 .

Thanks to the faster contraction rate $\beta\leq\gamma$, the advantage of the projected operator $\Pi_{\mathcal{Q}}\mathcal{R}^{\pi,\mu}$ is two-fold: (1) the operator often contracts faster to the limiting distribution $\eta_{\mathcal{R}}^{\pi}$ than the one-step operator $\mathcal{T}^{\pi}$ contracts to its own limiting distribution $\eta_{\mathcal{T}^{\pi}}$ (dabney2018distributional, ); (2) the limiting distribution $\eta_{\mathcal{R}}^{\pi}$ also enjoys a better approximation bound to the target distribution. We verify such results in Section 7.

Below, we use $z_{i}(x,a)$ to represent the $i$-th quantile of the distribution at $(x,a)$. Overall, we have a tabular quantile representation $\eta_{z}(x,a)=\frac{1}{m}\sum_{i=1}^{m}\delta_{z_{i}(x,a)},\forall(x,a)\in\mathcal{X}\times\mathcal{A}$, where we use the notation $\eta_{z}$ to stress the distribution’s dependency on parameter $z_{i}(x,a)$. For any given bootstrapping distribution vector $\eta\in\mathscr{P}_{\infty}(\mathbb{R})^{\mathcal{X}\times\mathcal{A}}$, in order to approximate the projected back-up target $\Pi_{\mathcal{Q}}\mathcal{R}^{\pi,\mu}\eta$ with the parameterized quantile distribution $\eta_{z}$, we solve the set of quantile regression problems for all $1\leq i\leq m,(x,a)\in\mathcal{X}\times\mathcal{A}$,

For any fixed $(x,a,i)$, to solve the quantile regression problem, we apply gradient descent on $z_{i}(x,a)$. In practice, with one sampled trajectory $(X_{s},A_{s},R_{s})_{s=0}^{\infty}\sim\mu$, the aim is to construct an unbiased stochastic gradient estimate of the QR loss $L_{z_{i}(x,a)}^{\tau_{i}}\left(\mathcal{R}^{\pi,\mu}\eta(x,a)\right)$. Below, let $\textrm{b}_{t}=\textrm{b}_{G_{0:t-1},\gamma^{t}}$ for simplicity. We start with a stochastic estimate $\widehat{L}_{z_{i}(x,a)}^{\tau_{i}}(\mathcal{R}^{\pi,\mu}\eta(x,a))$ for the QR loss,

Since $\widehat{L}_{z_{i}(x,a)}^{\tau_{i}}(\mathcal{R}^{\pi,\mu}\eta(x,a))$ is differentiable with $z_{i}(x,a)$, we use $\nabla_{z_{i}(x,a)}\widehat{L}_{z_{i}(x,a)}^{\tau_{i}}(\mathcal{R}^{\pi,\mu}\eta(x,a))$ as the stochastic gradient estimate. This gradient estimate is unbiased under mild conditions.

The above stochastic estimate bypasses the challenge that the QR loss is only defined against distributions, whereas sampled back-up targets $\widehat{R}^{\pi,\mu}\eta(x,a)=\eta(x,a)+\sum_{t=0}^{\infty}c_{1:t}\widetilde{\Delta}_{0:t}^{\pi}$ are signed measures in general. In Quantile Regression-Retrace, we use $\eta_{z}$ itself as the bootstrapping distribution, such that the algorithm approximates the fixed point iteration $\eta_{z}\leftarrow\Pi_{\mathcal{Q}}\mathcal{R}^{\pi,\mu}\eta_{z}$. Concretely, we carry out the following sample-based update

We introduce a deep RL implementation of the Quantile Regression-Retrace: QR-DQN-Retrace, where the parametric representation is combined with function approximations (bellemare2017cramer, ; dabney2018distributional, ; dabney2018implicit, ). The base agent QR-DQN (bellemare2017cramer, ) parameterizes the quantile locations $z_{i}(x,a;w)$ with the output of a neural network with weights $w$. Let $\eta(x,a;w)=\frac{1}{m}\sum_{i=1}^{m}\delta_{z_{i}(x,a;w)}$ denote the parameterized distribution. QR-DQN-Retrace updates its parameters by stochastic gradient descent on the estimated QR loss, averaged across all $m$ quantile levels $w\leftarrow w-\alpha\frac{1}{m}\sum_{i=1}^{m}\nabla_{w}\widehat{L}_{z_{i}(x,a;w)}^{\tau_{i}}\left(\mathcal{R}^{\pi,\mu}\eta(x,a;w)\right)$. In practice, the update is further averaged over state-action pairs sampled from a replay buffer.

We verify a few important properties of the distributional Retrace operator on a tabular MDP. The results corroborate the theoretical results from previous sections. Throughout, we use quantile representations with $m=100$ atoms; we obtain similar results for categorical representations. See Appendix F for details on the experiment setup. Let $\eta_{0}$ be the initial distribution, we carry out dynamic programming with $\mathcal{R}^{\pi,\mu}$ and denote $\eta_{k}=(\mathcal{R}^{\pi,\mu})^{k}\eta_{0}$ as the $k$^th distribution iterate.

We consider the control setting where the target policy $\pi$ is the greedy policy with respect to the Q-function induced by the parameterized distribution. Because the training data is sampled from a replay, the behavior policy $\mu$ is $\varepsilon$-greedy with respect to Q-functions induced by previous copies of the parameterized distribution. We evaluate the performance of deep RL agents on 57 Atari games (bellemare2013arcade, ). To ensure fair comparison across games, we compute the human normalized scores for each agent, and compare their evaluated mean and median scores across all 57 games during training.