A Finite Sample Analysis of Distributional TD Learning with Linear Function Approximation

Thank you for your encouraging comments and constructive suggestions. We are very glad to know that you find our theoretical results (Linear-CTD proposed in the paper achieves finite-sample convergence rates that are consistent with classical TD rates, and are independent of the number of categorical supports $K$, a key improvement over vanilla approaches) novel, and agree that the key conclusion of our paper (learning the full return distribution is no harder than learning its expectation with function approximation) is an important insight in RL. We are also grateful for your positive feedback on our writing (particularly Section 5) and the experimental section.

Regarding your concerns, we provide the following responses:

1. One suggestion for improvement would be to include a table summarizing the main theoretical results and comparing them with prior work, which could enhance clarity and accessibility (the current version is somewhat dense and requires flipping back and forth to see the comparison).

Thanks for your suggestion of presenting a table summarizing the main theoretical results and comparing them with prior work, we agree with you that it will enhance clarity and accessibility.

We will add the following table in the revised version. In the table, when the task is policy evaluation, the sample complexity is defined in terms of the $\mu_\pi$-weighted $L^2$ norm as the measure of error; when the task is distributional policy evaluation, the sample complexity is defined in terms of the $\mu_\pi$-weighted $W_1$ metric as the measure of error. This allows for a clearer comparison of theoretical results with prior work.

Here, Vanilla Linear-CTD refers to the stochastic semi-gradient descent (SSGD) algorithm with the probability mass function (PMF) representational (see Eqn.(26) in Appendix D.7). The contraction analysis in [5] in the table means that they provided a contraction analysis for the dynamic programming version of Vanilla Linear-CTD algorithm.

2. Additionally, the authors may consider presenting some of the empirical results in the main text, despite space limitations.

Thank you for your suggestion. We agree that presenting some experimental results in the main text would be helpful for readers. We will move Figures 1 and 2 to the main text in the revised version. Figure 1 demonstrates the convergence results of our algorithm. By comparing Figure 1 with Figure 2, we can verify the advantage of our Linear-CTD over the baseline algorithm (Vanilla Linear-CTD) as mentioned in Remark 5: when $K$ increases, the step size of the baseline algorithm needs to decay (Lemma E.3). In contrast, the step size of our Linear-CTD does not require adjustment as $K$ increases (Lemma 5.3).

[1] Gen Li, Weichen Wu, Yuejie Chi, Cong Ma, Alessandro Rinaldo, and Yuting Wei. High-probability sample complexities for policy evaluation with linear function approximation. IEEE Transactions on Information Theory, 2024.

[2] Sergey Samsonov, Daniil Tiapkin, Alexey Naumov, and Eric Moulines. Improved high-probability bounds for the temporal difference learning algorithm via exponential stability. In The Thirty Seventh Annual Conference on Learning Theory, pages 4511–4547. PMLR, 2024.

[3] Weichen Wu, Gen Li, Yuting Wei, and Alessandro Rinaldo. Statistical Inference for Temporal Difference Learning with Linear Function Approximation. arXiv preprint arXiv:2410.16106, 2024.

[4] Tianjiao Li, Guanghui Lan, and Ashwin Pananjady. Accelerated and instance-optimal policy evaluation with linear function approximation. SIAM Journal on Mathematics of Data Science, 5(1):174–200, 2023.

[5] Marc G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. MIT Press, 2023. http://www.distributional-rl.org.

We sincerely appreciate your effort and time in reviewing our paper. And we are glad to hear that you find our results are novel and technically sound. Below, we hope our responses can address your concerns.

Weakness 1. The paper's contributions are heavily reliant on pre-existing results, which dampens potential originality. Most of the basic technical challenges have been considered in prior work, such as those relating to Wasserstein-metric convergence or analysis under non-i.i.d. sampling.

While we have drawn upon some pre-existing results, including the Wasserstein-metric convergence of the baseline algorithm and the theoretical framework under non-i.i.d. sampling, we maintain that our work has significant technical originality and contributions.

1. Regarding the Wasserstein-metric convergence you mentioned, we speculate that you are referring to the contraction analysis for the dynamic programming version of the baseline algorithm Vanilla Linear-CTD provided by Section 9.7 in [1]. To the best of our knowledge, there is a substantial gap between the contraction analysis of dynamic programming algorithms and the finite-sample analysis of online algorithm (e.g. from iterative policy evaluation (dynamic programming) to TD learning, and from value iteration to Q-Learning).

2. While we employed the theoretical framework of linear stochastic approximation from [2] (encompassing both i.i.d. and non-i.i.d. scenarios) to analyze the sample complexity, significant efforts were required to derive bounds that are independent of $K$ and match those of classic Linear-TD. These include a novel analysis of the upper bound on the spectral norm of matrix $C\tilde G^{-1}(r)C^{-1}$ in Section 5.3, an upper bound on the spectral norm of the expectation of the Kronecker product presented in Lemma 5.4, and the technique of converting the bounding of the norm of $b_t$ into bounding a $\ell_2$ distance between two probability distributions. We believe these results are of independent interest and will be beneficial for subsequent research in distributional reinforcement learning.

We also would like to mention our algorithmic contributions, which is fully original. To be specific, we proposed an improved version of the linear-categorical TD learning algorithm (Linear-CTD). Compared to the baseline algorithm, which we refer to as Vanilla Linear-CTD, our algorithm is a preconditioned variant that is more robust, as its step size is independent of the number of support points $K$ (as noted by Reviewer PXyE), which is supported by extensive numerical experiments.

In summary, we believe that our work demonstrates significant originality and contributions compared to pre-existing results. We hope that the response clarifies this point.

Weakness 2. However, there is no compelling need shown in this paper on why an in-depth analysis of distributional convergence is necessary.

We would believe in-depth theoretical analyses of distribution policy evaluation algorithms is timely and important. In the original version, we reviewed existing theoretical work on TD learning [2, 3, 4, 5], distributional TD learning in the tabular setting [6, 7, 8, 9, 10], as well as asymptotic analyses of distributional TD learning under linear function approximation [1]. Building on these prior studies, we argued that investigating the non-asymptotic convergence rate of distributional TD learning under linear function approximation is a natural and meaningful extension.

In the revised version, we will include the following content in the Introduction section to further clarify the significance of an in-depth analysis of distributional convergence.

1. In risk-sensitive domains such as finance [11] and healthcare [12], decision-making relies not only on the mean of returns but also on some risk metrics (e.g., variance, quantile, CVaR). Distributional convergence results offer theoretical guarantees for the reliable estimation of these risk metrics in practical tasks.

2. Many risk-sensitive RL algorithms [13, 14] incorporate distributional RL as a core step to estimate specific expected utility functions of return distributions. The distributional convergence results can provide finite-sample convergence rates for this critical step, giving theoretical guarantees for these algorithms.

For further discussions on the significance of studying distributional RL algorithms, it could be also referred to Chapter 1 in [1] and Section 2.2 in [15].

[1] Marc G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. MIT Press, 2023. http://www.distributional-rl.org.

[2] Sergey Samsonov, Daniil Tiapkin, Alexey Naumov, and Eric Moulines. Improved high-probability bounds for the temporal difference learning algorithm via exponential stability. In The Thirty Seventh Annual Conference on Learning Theory, pages 4511–4547. PMLR, 2024.

[3] Gen Li, Changxiao Cai, Yuxin Chen, Yuting Wei, and Yuejie Chi. Is q-learning minimax optimal? a tight sample complexity analysis. Operations Research, 72(1):222–236, 2024a.

[4] Gen Li, Weichen Wu, Yuejie Chi, Cong Ma, Alessandro Rinaldo, and Yuting Wei. High-probability sample complexities for policy evaluation with linear function approximation. IEEE Transactions on Information Theory, 2024.

[5] Weichen Wu, Gen Li, Yuting Wei, and Alessandro Rinaldo. Statistical Inference for Temporal Difference Learning with Linear Function Approximation. arXiv preprint arXiv:2410.16106, 2024.

[6] Mark Rowland, Marc Bellemare, Will Dabney, R´emi Munos, and Yee Whye Teh. An analysis of categorical distributional reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pages 29–37. PMLR, 2018.

[7] Mark Rowland, R´emi 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. Journal of Machine Learning Research, 25:1–47, 2024.

[8] Liangyu Zhang, Yang Peng, Jiadong Liang, Wenhao Yang, and Zhihua Zhang. Estimation and Inference in Distributional Reinforcement Learning. The Annals of Statistics, 2025. Forthcoming.

[9] Mark Rowland, Li Kevin Wenliang, Remi Munos, Clare Lyle, Yunhao Tang, and Will Dabney. Near-minimax-optimal distributional reinforcement learning with a generative model. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024

[10] Yang Peng, Liangyu Zhang, and Zhihua Zhang. Statistical efficiency of distributional temporal difference learning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.

[11] Haoran Wang and Xun Yu Zhou. Continuous-time mean–variance portfolio selection: A reinforcement learning framework. Mathematical Finance, 30(4): 1273–1308, 2020.

[12] Lan Wang, Yu Zhou, Rui Song, and Ben Sherwood. Quantile-optimal treatment regimes. Journal of the American Statistical Association, 113(523):1243–1254,2018.

[13] Mehrdad Moghimi and Hyejin Ku. Beyond CVar: Leveraging static spectral risk measures for enhanced decision-making in distributional reinforcement learning. In Forty-second International Conference on Machine Learning, 2025. URL https://openreview.net/forum?id=WeMpvGxXMn.

[14] Bernardo Avila Pires, Mark Rowland, Diana Borsa, Zhaohan Daniel Guo, Khimya Khetarpal, Andre Barreto, David Abel, Remi Munos, and Will Dabney. Optimizing return distributions with distributional dynamic programming. arXiv preprint arXiv:2501.13028, 2025.

[15] Zhengling Qi, Chenjia Bai, Zhaoran Wang, and Lan Wang. Distributional off-policy evaluation in reinforcement learning. Journal of the American Statistical Association, jun 2025. doi: 10.1080/01621459.2025.2506197. Forthcoming.

Thank you for your encouraging comments. We are very glad to know that you think our paper well-organized and readable, and our Linear-CTD robust and novel compared with baseline method. Regarding the weaknesses, minor comments and questions, we provide the following detailed responses:

Weakness: Although some explanation is provided in the appendix, the meaning and implication of Theorem 3.1 may not be immediately clear to readers unfamiliar with this line of work. For better readability, a short paragraph providing intuition or commentary around the theorem in the main text would be helpful.

Thank you for your feedback. We fully agree that adding more explanations for the linear-categorical projected Bellman equation (Theorem 3.1) will enhance readability. In the original version, due to space constraints some critical intermediate results had to be placed in the appendix, such as Proposition D.3, which characterizes the properties of the categorical projected Bellman operator and introduces the notations $g$ and $G$ for the first time. In the revised version, we will move this proposition back to the main text and systematically elaborate on the important concepts including $g$ and $G$ which are related to the categorical projected Bellman operator. To further help readers understand, we will add the following table in the revised version that compares various concepts between Linear-TD and Linear-CTD. We believe this will improve the readability of the paper.

Minor Comments

Thank you for your reminder. The important works [1, 2] are not limited to the classic categorical and quantile representations; instead, they discuss distributional RL algorithms that consider general statistical functional representations. We will discuss these works in the revised version.

Questions: Could the author elaborate on the meaning of Equation (10) and the footnote that follows it? Specifically, is the term $(k+1)/(K+1)$ used as a technical convenience for representing the discrete uniform distribution in terms of $1_K^\top$? If so, does it also help improve the tightness of the derived sample complexity bounds?

Thank you for your valuable questions and suggestions. The term $(k+1)/(K+1)$ in Eqn. (10) is indeed the cdf of discrete uniform distribution $\nu$ on $\\{x\_0,\ldots,x\_K\\}$ at $x_k$. And as you noted, introducing normalization indeed helps improve the tightness of the derived sample complexity bound and is indispensable. In addition, the normalization also offers algorithmic and interpretability benefits. We will elaborate on normalization in detail below and add this content in the revised version. We believe that this will enhance readers' understanding.

As stated in footnote 2, the $(k+1)/(K+1)$ term corresponds to the cdf of a discrete uniform distribution $\nu$ on $\\{x\_0,\ldots,x\_K\\}$ at $x_k$ for normalization, guaranteeing our estimator has mass $1$. First, we explain why normalization is necessary. Although we work with signed measures, it remains necessary to ensure the total mass is normalized to $1$. The reasons are from both algorithmic and theoretical considerations. From an algorithmic perspective, without normalization, the difference between our estimator and the true return distribution would not be a signed measure with mass zero. This would make our estimator less interpretable and lead to issues in defining many important quantities. For instance, the $W_1$ distance would lose its integral representation $W_1(\nu_1,\nu_2)=\sup_{f:Lip(f)\leq 1}(\nu_1-\nu_2)f$. Furthermore, it would fail to generate the same iterative sequence as classic TD (Proposition D.4) and would not be equivalent to tabular CTD in the tabular case. In this sense, the algorithm will not be a natural extension of Linear-TD and tabular CTD. From a theoretical perspective, tight sample complexity bounds can only be achieved after introducing normalization. Without normalization, the $1_K$ term would be absent from Eqn. (18), which would lead to a worse bound of $\|\|b_t\|\|$, and furthermore, the sample complexity of Linear-CTD cannot match that of classic Linear-TD.

Next, we explain the impact of choosing different normalizers (distributions used for normalization) and why we choose the discrete uniform distribution $\nu$ as the normalizer. As noted in footnote 2, the definition of space $\mathscr{P}_{\phi,K}^{sign}$ depends on the choice of the normalizer, further affecting the optimal solution and the iterating process of the algorithm, an outcome we aim to avoid. Fortunately, such undesirable properties would not appear in the tabular case. In the tabular case, it can be easily shown that the Linear-CTD algorithm is equivalent to tabular CTD regardless of the distribution chosen for normalization. In the general case, we can introduce an intercept term (a state-independent feature that is always $1$) in the feature $\phi$ to ensure that the choice of the distribution has no impact on the algorithm (because we can always adjust the coefficients of the intercept terms to offset the impact of changing the normalizers). Based on the above discussion, any distribution could be used for normalization in practice.

And we choose the uniform distribution as the normalizer only for conceptual and technical simplicity. For example, in Section 5.3 Bounding the Norm of $b_t$ paragraph, using the discrete uniform distribution $\nu$ allows us to more conveniently derive results like Lemma G.4, thereby making it easy to get a tight bound for $\|\|b_t\|\|$. We hope this explanation is helpful.

[1] Rowland, Mark, et al. Statistics and samples in distributional reinforcement learning. International Conference on Machine Learning. PMLR, 2019.

[2] Cho, Taehyun, et al. Tractable and Provably Efficient Distributional Reinforcement Learning with General Value Function Approximation. arXiv preprint arXiv:2407.21260 (2024).

Thank you for your valuable review and suggestions. We are very glad to know that you think our paper is well-written, and our results are helpful for subsequent research in distributional RL.

Regarding the weaknesses and questions, we provide the following detailed responses:

Weakness 1: The proposed theory and algorithms apply to the space of signed measures, rather than usual unsigned measures. And Question 2: What are the limitations of working with signed measures?

We are grateful to you for raising this issue. We would like to point out similar issues also appear in works such as [1, 2] and Section 9.6 in [3], and our new numerical experiments indicate that in non-tabular settings, non-probability signed measures generally arise. The non-probability signed measure does induce some limitations. For instance, this makes some common statistical functionals such as quantiles and CVaR ill-defined, thereby failing to guarantee their convergence. We will provide a more detailed discussion in the revised version.

However, according to our Theorem 4.1, for statistical functionals that are Lipschitz continuous with respect to the $W_1$ distance (e.g., expected utility when the utility function is Lipschitz), we can still establish their finite-sample convergence. Thus, our finite-sample theory built on signed measure spaces still finds applications in many problems. For example, some risk-sensitive RL algorithms, such as [4, 5], use distributional RL algorithms to estimate specific expected utility functions of the distributions, and our results can give a finite sample convergence rate.

To avoid issues caused by non-probability signed measures, one could consider linear function approximation with softmax-parameterized categorical representations, or quantile representations in Section 9.5 in [3]. However, these parameterizations leads to nonlinear algorithms, which may prevent us from proving finite-sample convergence comparable to that of Linear-TD. Compared to these two nonlinear algorithms, the introduction of signed measures in our Linear-CTD and other works helps preserve the linear structure of the algorithm, enabling us to obtain convergence results that match those of Linear-TD. Regarding the finite-sample analysis of the aforementioned nonlinear algorithms, we leave this as future work. We hope this explanation is helpful.

Question 1: Why have you added $(k+1)/(K+1)$ in (10)? Please add detail on your claim that it is for normalization.

Thank you for your valuable questions and suggestions. This term serves as a normalizer. Below we will provide a detailed explanation and include it in the revised version.

As stated in footnote 2, the $(k+1)/(K+1)$ term corresponds to the cdf of a discrete uniform distribution $\nu$ on $\\{x\_0,\ldots,x\_K\\}$ at $x_k$ for normalization, guaranteeing our estimator has mass $1$. First, we explain why normalization is necessary. Although we work with signed measures, it remains necessary to ensure the total mass is normalized to $1$. The reasons are from both algorithmic and theoretical considerations. From an algorithmic perspective, without normalization, the difference between our estimator and the true return distribution would not be a signed measure with mass zero. This would make our estimator less interpretable and lead to issues in defining many important quantities. For instance, the $W_1$ distance would lose its integral representation $W_1(\nu_1,\nu_2)=\sup_{f:Lip(f)\leq 1}(\nu_1-\nu_2)f$. Furthermore, it would fail to generate the same iterative sequence as classic TD (Proposition D.4) and would not be equivalent to tabular CTD in the tabular case. In this sense, the algorithm will not be a natural extension of Linear-TD and tabular CTD. From a theoretical perspective, tight sample complexity bounds can only be achieved after introducing normalization. Without normalization, the $1_K$ term would be absent from Eqn. (18), which would lead to a worse bound of $\|\|b_t\|\|$, and furthermore, the sample complexity of Linear-CTD cannot match that of classic Linear-TD.

Next, we explain the impact of choosing different normalizers (distributions used for normalization) and why we choose the discrete uniform distribution $\nu$ as the normalizer. As noted in footnote 2, the definition of space $\mathscr{P}_{\phi,K}^{sign}$ depends on the choice of the normalizer, further affecting the optimal solution and the iterating process of the algorithm, an outcome we aim to avoid. Fortunately, such undesirable properties would not appear in the tabular case. In the tabular case, it can be easily shown that the Linear-CTD algorithm is equivalent to tabular CTD regardless of the distribution chosen for normalization. In the general case, we can introduce an intercept term (a state-independent feature that is always $1$) in the feature $\phi$ to ensure that the choice of the distribution has no impact on the algorithm (because we can always adjust the coefficients of the intercept terms to offset the impact of changing the normalizers). Based on the above discussion, any distribution could be used for normalization in practice.

And we choose the uniform distribution as the normalizer only for conceptual and technical simplicity. For example, in Section 5.3 Bounding the Norm of $b_t$ paragraph, using the discrete uniform distribution $\nu$ allows us to more conveniently derive results like Lemma G.4, thereby making it easy to get a tight bound for $\|\|b_t\|\|$. We hope this explanation is helpful.

Question 2 Cont.: Suppose we use a notion of ordering between the return distributions of different policies, e.g., first order stochastic dominance. Would this ordering continue to hold when we project them onto the space of signed measures?

Thank you for raising this valuable question. Through theoretical analysis, we find that the answer is no. However, the reason why the ordering no longer holds generally is not due to the issue of non-probability signed measures, but the linear function approximation. Here we give the justification: Let $\xi_1,\xi_2\in\mathscr{P}^{\mathcal{S}}$ such that $\xi_1$ stochastic dominant $\xi_2$, that is $F_{\xi_1(s)}(x)\geq F_{\xi_2(s)}(x)$ for any $s\in\mathcal{S}$ and $x\in\mathbb{R}$. We denote by $\mathbf{F}\_{\xi_i}(s)=(F\_{\xi_i(s)}(x_k))\_{k\in\{0,1,\ldots,K-1\}}$ the cdf vector of the categorical projection $\Pi\_K \xi\_i(s)\in\mathscr{P}\_K$. Then using $\Pi\_K$ is a monotone map (Proposition 5 in [6]), we have $\mathbf{F}\_{\xi\_1}(s)\geq\mathbf{F}\_{\xi\_2}(s)$ element-wisely. By Proposition 3.1, the difference of the cdf vectors of $(\Pi^\pi\_{\phi, K}\xi\_1)(s^\prime),(\Pi^\pi\_{\phi, K}\xi\_2)(s^\prime)\in\mathscr{P}\_K$ is given by $\tilde{\mathbf F}\_{\xi\_1}(s^\prime)-\tilde{\mathbf F}\_{\xi\_2}(s^\prime)=\mathbb{E}\_{s\sim\mu\_{\pi}}[(\mathbf{F}\_{\xi\_1}(s)-\mathbf{F}\_{\xi_2}(s))\phi(s)^{\top}]\Sigma\_{\phi}^{-1}\phi(s^\prime)$. In the tabular case, it is easy to see that $\tilde{\mathbf F}\_{\xi_1}(s^\prime)\geq\tilde{\mathbf F}\_{\xi_2}(s^\prime)$, but in the linear approximation case, it is not necessarily non-negative because $\phi(s)^{\top}\Sigma\_{\phi}^{-1}\phi(s^\prime)$ can be negative and one can construct a counterexample based on the fact. Moreover, if we consider the more general case, namely we apply two different projection operators of different policies to $\xi\_1$ and $\xi\_2$ respectively, the answer is still no.

In the simple classic Linear-TD, the same argument holds: Let $V_1, V_2\in\mathbb{R}^{\mathcal{S}}$ such that $V_1\geq V_2$, then we have $(\Pi^\pi\_{\phi}V_1)(s^\prime)-(\Pi^\pi\_{\phi}V_2)(s^\prime)=\mathbb{E}\_{s\sim\mu_{\pi}}[(V_1(s)-V_2(s))\phi(s)^{\top}]\Sigma\_{\phi}^{-1}\phi(s^\prime)$, which is not necessarily non-negative, except in the tabular case.

Question 3: More generally, policy iteration with function approximation is known to suffer from policy oscillation, divergence and other issues (Bertsekas, 2011). Do you think Policy Iteration with distributional policy evaluation or some variant address those issues?

Thank you for this highly valuable question. In the case of linear approximation, we believe the answer is no. As stated in Proposition D.4 of our paper, under linear function approximation, distributional RL algorithms generate an identical sequence of value function estimates to classic RL algorithms. Thus, introducing distribution estimation cannot alleviate these problems.

In the general nonlinear function approximation setting, as demonstrated in many empirically focused studies, distributional RL may empirically mitigate such problems, which is likely due to enhanced representation learning enabled by richer learning objectives.

We hope this response is helpful.

[1] Marc G Bellemare, Nicolas Le Roux, Pablo Samuel Castro, and Subhodeep Moitra. Distributional reinforcement learning with linear function approximation. AISTATS 2019.

[2] Clare Lyle, Marc G Bellemare, and Pablo Samuel Castro. A comparative analysis of expected and distributional reinforcement learning. AAAI 2019.

[3] Marc G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. MIT Press, 2023. http://www.distributional-rl.org.

[4] Mehrdad Moghimi and Hyejin Ku. Beyond CVar: Leveraging static spectral risk measures for enhanced decision-making in distributional reinforcement learning. ICML 2025.

[5] Bernardo Avila Pires, Mark Rowland, Diana Borsa, Zhaohan Daniel Guo, Khimya Khetarpal, Andre Barreto, David Abel, Remi Munos, and Will Dabney. Optimizing return distributions with distributional dynamic programming. arXiv preprint arXiv:2501.13028, 2025.

[6] Mark Rowland, Marc Bellemare, Will Dabney, R´emi Munos, and Yee Whye Teh. An analysis of categorical distributional reinforcement learning. AISTATS 2018.