A Finite Sample Analysis of Distributional TD Learning with Linear Function Approximation
Authors: Yang Peng
Kaicheng Jin
Liangyu Zhang
Zhihua Zhang
Venue: NeurIPS 2025 poster ()
Keywords: Distributional Reinforcement Learning
Sample Complexity
Distributional TD Learning
Distributional Policy Evaluation
Stochastic Approximation
Primary area: reinforcement_learning

Abstract: In this paper, we study the finite-sample statistical rates of distributional temporal difference (TD) learning with linear function approximation.
The aim of distributional TD learning is to estimate the return distribution of a discounted Markov decision process for a given policy $\pi$.
Previous works on statistical analysis of distributional TD learning mainly focus on the tabular case.
In contrast, we first consider the linear function approximation setting and derive sharp finite-sample rates.
Our theoretical results demonstrate that the sample complexity of linear distributional TD learning matches that of classic linear TD learning.
This implies that, with linear function approximation, learning the full distribution of the return from streaming data is no more difficult than learning its expectation (value function).
To derive tight sample complexity bounds, we conduct a fine-grained analysis of the linear-categorical Bellman equation and employ the exponential stability arguments for products of random matrices.
Our results provide new insights into the statistical efficiency of distributional reinforcement learning algorithms.

A Finite Sample Analysis of Distributional TD Learning with Linear Function Approximation
Yang Peng"
Kaicheng Jin
Liangyu Zhang
Zhihua Zhang
Abstract
In this paper, study the finite-sample statistical rates of distributional temporal difference (TD) learning with linear function approximation. The aim of distribu- tional TD learning is to estimate the return distribution of discounted Markov decision process for a given policy т. Previous works on statistical analysis distributional TD learning mainly focus the tabular case. In contrast, we first consider the linear function approximation setting and derive sharp finite-sample rates. Our theoretical results demonstrate that the sample complexity of linear distributional TD learning matches that of classic linear TD learning. This implies that, with linear function approximation, learning the full distribution of the return from streaming data is no more difficult than learning its expectation (value func- tion). To derive tight sample complexity bounds, conduct fine-grained analysis of the linear-categorical Bellman equation and employ the exponential stability arguments for products of random matrices. Our results provide new insights into the statistical efficiency of distributional reinforcement learning algorithms.
1 Introduction
Distributional policy evaluation [Morimura al., 2010, Bellemare et al., 2017, 2023], which aims to estimate the return distribution of policy in an Markov decision process (MDP), is crucial for many uncertainty-aware or risk-sensitive tasks [Lim and Malik, 2022, Kastner et al., 2023]. Unlike the classic policy evaluation that only focuses on expected returns (value functions), distributional policy evaluation captures uncertainty and risk by considering the full distributional information. To solve distributional policy evaluation problem, in the seminal work Bellemare al. [2017] proposed distributional temporal difference (TD) learning, which can be viewed as an extension of classic TD learning [Sutton, 1988].
Although classic TD learning has been extensively studied [Bertsekas and Tsitsiklis, 1995, Tsitsiklis and Van Roy, 1996, Bhandari al., 2018, Dalal et al., 2018, Patil et al., 2023, Li et al., 2024a,b, Chen al., 2024, Samsonov et al., b024a. Wu al., 2024], the theoretical understanding of distributional TD learning, which is important for risk- tasks [Wang and Zhou, 2020, Wang et al., 2018, Moghimi and Ku, 2025, Avila Pires et al.. 2025, Qi et al., 2025]. remains relatively underdeveloped. Recent works [Rowland al., 2018, Böck and Heitzinger, 2022, Zhang et al., 2025, Rowland et al., b024a,b, Peng et al., 2024] have analyzed distributional TD learning (or its model-hased variants) in the tabular setting. Especially, Rowland al. [2024b] and Peng et al. [2024] demonstrated that the
School of Mathematical Sciences, Peking University; email: pengyangēpku. edu. School of Mathematical Sciences, Peking University: email: kcjin0pku. edu. cn. 1 School of Statistics and Data Science, Shanghai University of Finance and Economies; email: zhangliangyu@sufee www.ssuuueeuuueee.. en.
$School of Mathematical Sciences, School of Computer Science, Peking University; Center for Intelligent Computing, Great Bay University; email: zhzhang@math, pku. edu.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

tabular setting, learning the return distribution (in terms of the 1- 1-Wasserstein distance") is statistically as easy as learning its expectation. However, in practical scenarios, where the state space is extremely large or continuous, the function approximation [Dabney et al., 201 8b,a, Rowland et al., 2019, Yang et al., 2019, Freirich et al., 2019, Yue et al., 2020, Nguyen-Tani TTng et al., 2021, Zhou et al., 2021, Luo et al., 2022, Wenliang et al., 2024, Sun et al., 2024, Cho et al., 2024, Shen et al., 2025] becomes indispensable. This raises a new open question: When function approximation is employed, does learning the return distribution remain statistically efficient as learning its expectation?
To answer this question, we consider the simplest form of function approximation, i.e., linear function approximation, and investigate the finite-sample performance of linear distributional TD learning. In distributional TD learning, we need to represent the tiliiiteddiensionaal return distributions with some mnite-dimensionll parametrizations to make the algorithm tractable. Previous works [Bellemare et al., 2019, Lyle et al., 2019, Bellemare et al., 2023] have proposed various línear distributional TD learning algorithms under different parameterizations, namely categorical and quantile parametrizations. In this paper, we consider the categorical parametrization and propose an improved version of the linear-categorical TD learning algorithm (Linear-CTD) We then analyze the non-asymptotic convergence rate of Linear-CID. Our analysis reveals that, with the Dolyak-Ruppens tail averaging [Ruppert, 1988, Polyak and Juditsky, 1992] and proper constant step size, the sample complexity of Linear-CiD CTD matches that of classic linear TD learning (Linear-TD) [Li et al., 2024b, Samsonov et al., 2024b]. Thus. this confirms that learning the return distribution is statistically no more difficult than learning its expectation when the linear function approximation is employed.
Notation. In the following parts of the paper. (x)+ == max {x, 0} for any I E R. "S" (resp. "2") means no larger (resp. smaller) than up to a multiplicatiie universal constant, and a ~ means a 2 b and a Z hold simultaneously. The asymptotic notation f(.) = Ô (g(+)) (resp. & (g(()) means that f(.) is ouder-wise no larger (resp. smaller) than g(.), ignoring logarithmic factors of polynomials of (1 (1-7)-1.Xp)) Am a- e 8-1, K, || Iz, |0= x We will explain the concrete meaning of the notation once we have encountered them for the first time.
We denote by öI the Dirac measure at x E IR, I the indicator function, & the Kronecker product (see Appendix A), &RR the all-ones vector, 0x C the all-zeros vector, E the identity matrix, ||w/|| the Euclidean norm of any vector u, []] the spectral norm of any matrix B, and ||ul| := VuT Bu when B is positive semi-definite (PSD). B stands for B for any symmetric matrices Bi, B2. And II B₁ is defined as B,B.- B, for any matrices (B4)k= with appropriate sizes. For any matrix B= b( we define its vector- ization as vec (B)=(b(1)T ..... ...., b(n)T)
B₁ is PSD
ER'm. TTRRmm Given set A, we denote by (AA) the set of all probability distributions over A. For simplicity, we abbreviate A[0 (1- as P.
Contributions. Our contribution is two-fold: in algorithms and in theory. Algorithmically, propose an improved version of the linear-categorical TD learning algorithm (Linear-CTD). Rather than using stochastic semi-gradient descent to update the parameter as in Bellemare et al. [2019], Lyle et al. [2019), Bellemare et al. [2023], we directly formulate the linear-categorical projected Bellman equation into linear system and apply a linear stochastic approximation to solve it. The resulting Linear-CTD can be viewed as preconditioned version [Chen, 2005, Li, 2017] of the vanilla linear categorical TD learning algorithm proposed in Bellemare et al. [2023, Section 9.6]. By introducing a preconditioner, our Linear-CTV achieves a finite-sample rate independent of the number of supports K in the categorical parameterization, which the vanilla version cannot attain. We provide both theoretical and experimental evidence to demonstrate this advantage of our Linear-CTD.
Theoretically, establish the first non-asy non-asympturn asymptotic guarantees for distributional TD learning with the linear function approximation. Specifically, we show that in the generative model setting, with the Polyak-Ruppers tail averaging and constant step size, we need
T= 5 (E > (1 1
mia (1 >
i +1)
online interactions with the environment to ensure Linear-CTD yields a e-accurate estimator with high probability, when the error is measured by the /e--weettted 1 1-Wasserstein distance. We also
*Solving distributional policy evaluation e-accurately in the 1- I-Wasserstein distance sense is harder than solving classic policy evaluation e-accurately, as the absolute difference of value functions is always bounded by the 1- L-Wasserstein distance between return distributions.

---

extend the result to the Markovian setting. Our sample complexity bounds match those of the classic Linear-TD with constant step size [Li et al., 2024b, Samsonov al., 2024b]. confirming the same statistical tractability of distributional and classic aalue-based policy evaluations. To establish these theoretical results, we analyze the linear-categorical Bellman equation in detail and apply the exponential stability argument proposed in Samsonov et al. [2024b]. Our analysis of the linear- categorical Bellman equation lays the foundation for subsequent algorithmic and theoretical advances in distributional reinforcement learning with function approximation.
Organization. The remainder of this paper is organized as follows. In Section 2. we recup Linear-TD and tabular categorical TD learning. In Section 3, we introduce the linear-categorical parametrization, and use the cintar-ctteooiical projected Bellman equation to derive Linear- CTD. In Section 4, we employ the exponential stability arguments to analyze the statistical efficiency of Linear-CTD. The proof is outlined in Section 5. In Section 6, we conclude our work. See Appendix B for more related work. In Appendix F, compare various concepts and results between Linear-TD and Linear- CTD. In Appendix G, we empirically validate the convergence of Linear-CTD and compare it with prior algorithms through numerical experiments, confirming our theoretical findings. Details of the proof are given in the appendices.
2 Backgrounds
In this section, we recap the basics of policy evaluation and distributional policy evaluation tasks.
2.1 Policy Evaluation
A discounted MDP is defined by a 4-tuple M = A, We assume that the state space S and the action space A are both Polish spaces, namely complete separable metric spaces. P(--- s, a) is the joint distribution of reward and next state condition on (s, a) E S x A. We assume that all rewards are bounded random variables in [0, 1]. And y E (0. 1) the discount factor.
Given a policy :: S + A (AA) and an initial state 8q = 8 E S, a random trajectory {(st, a, rt) ] can be sampled: at | 8t T(- 8t). (rt, 8++1) (z) P(::- 8t, ax), for any E N. We assume the Markov chain (st) has a unique stationary distribution M4 E A(S). We define the return of the trajectory as (T((s == Ei= 7'r,. The value function (*()) is the expectation of GT (s), and
V= = -V* ERS. It is known that V satisfies the Bellman equation:
(s) = (()=
+s),, Vs E.
or in a compact form V V* where T RSS - R is called the Bellman operator. In the task of policy evaluation, we aim to find the unique solution V of the equation for some given policy π.
Tabular TD Learning. The policy evaluation problem reduced to solving the Bellman equation. However, in practical applications T" is usually unknown and the agent only has access to the streaming data {(s,, a,, F2) r=) In this cireumstance, we can solve the Bellman equation through linear stochastic approximation (LSA). Specifically, in the t-th time-step the updating scheme is
(s(s)) + Vs:1 (V:) ra -V-11()441), Ve(s) + v)-1(s), Vs + St-
(2)
We expect V. to converge to V as tends to infinity. This algorithm is known as TD learning, however, it is computationally tractable only in the tabular setting.
Linear Function Approximation and Linear TD Learning. In this part, we introduce linear function approximation and briefly review the more practical Linear-TD. To be concrete, assume there is a d-dimensional feature vector for each state 8 E S, which is given by the feature map Φb S + RRd. We consider the linear function approximation of value functions:
Y = VV+- = (Vx( Vv((s) (bs) T sb 2 E (R)) C R...
(3)
My M--weighted norm
V Nr (Esnjay 1/, and linear projection operator II: RS Yg:
II V argming EX Vøllpa V..RS.
3

---

One can check that the linear projected Bellman operator I is y-contraction in the Polish space (fp. 1-11pa> Hence, T admits unique fixed point V which satisfies (1--2) -// V lluu- [Bellemare et al., 2023, Theorem 9.8]. In Appendix C.I, we show that R2:
llp S
2h* _ is the unique solution to the linear system for xp
(4) In the subscript of the expectation, abbreviate Me ()), (--ss), (r, P(.. s, s, as 8, a, r, s'. For brevity, we will use such abbreviations in this paper when there is no ambiguity. We use LSA solve the linear projected Bellman equation (Eqn. 4). As result, at the t-th time-step, the updating scheme of Linear-TD is
- 7Ba, (ss))((s))] + Ee, [o(s)r), Eg :- :o [o(s))(s)
Linear-TD:
w t -h adb(s)
(db(t)) r0(s)+1)) wh-11 - m1).
-1
(5)
2.2 Distributional Policy Evaluation
In certain applications, we are not only interested in finding the expectation of random return G (s)
also want to find the whole distribution of (s). This task called distributional policy evaluation. We use (s) E S to denote the distribution of GT (s) and let n7 := (s))sEs) E go Then 17th satisfies the distributional Bellman equation:
17
"" ()")))
Vs E
(6)
where the RHS is the distribution of ro +o conditioned on so = 8. Here (x) : for any R. and f 3 is defined f A) == off=: f(x:) A}) for any function f: IR R,
probability measure V E 30 and Borel set C R. The distributional Bellman equation can also be written 17" T 17. The operator T 9o5 - gos is called the distributional Bellman operator.
In this task, our goal is to find 17t for some given policy т.
Tabular Distributional TD Learning. In analogy tabular TD learning (Eqn. (2)), in the tabular setting, we can solve the distributional Bellman equation by LSA and derive the distributional TD learning rule given the streaming data (e)),
(s:) + ( - a[n-1(()) (br..) # n ( me(s) + m)-1(s), Vs + $t-
We comment the algorithm above is not computationally feasible need to manipulate infinite- dimensional objects (return distributions) at each iteration.
Categorical Parametrization and Tabular Categorical TD Learning. In order deal with return distributions in computationally tractable manner, we consider the categorical parametrization as in Bellemare et al. [2017], Rowland et al. [2018, 2024b], Peng et al. [2024]. To be compatible with linear function approximation introduced in the next section, which cannot guarantee non-negative outputs, will work with gosign the signed measure space with total mass 1 as in Bellemare et al. [2019], Lyle et al. [2019], Bellemare et al. [2023] instead of standard probability space 3 C gosign
g = v: v [ = 1, supp(v) C (1-
For any E Gesign define its cumulative distribution function (CDF) F2(x) v(10, æl). We can naturally define the L and L' distances between CDFs as the Cramér distance & and 1- L-Wasserstein distance W1 in gesign respectively. The distributional Bellman operator (see Eqn. (6)) can also be extended to the product space (gasinns without modifying its definition.
The space of all categorical parametrized signed measures with total mass 1 is defined as
= {up = Ek Próua p (Po. .PD.....PPE. .K- K2 PK 1- pk).
K K
(7)
which is an affine subspace of @sign, Here { k= are K+1 equally-spaced points of the support, <= [K(1~))] is the gap between adjacent points, and pa the 'probability' (may be
negative) that assigns to Fg- We define the categorical projection operator Ix gsign Signn as
Il v == argmin
PE M0 l2 (v,up), Yu E grign

---

Following Bellemare et al. [2023, Proposition 5.14], one can show that H v Pxx is uniquely represented with vector p. (px(v)) _2 where
(a(1) Jo.a-) - (x - -xx/////// =x))/k((+(dd))))
(8)
We lift II to the product space by defining [I (s) : Ix One can check that the categorical Bellman operator II T is VCcontractionn in the Polish space (( p2,
where 2, (m, := (Es-u,, (2))) is the μe, uu-weighted Cramér distance between m, 72 E (S Similarly, WT,,, (m, 772) (Bx-ne, (Bs= (m) (s), (2()))))22 Hence, the cate- gorical projected Bellman equation 7 1 T admits unique solution n7 which satisfies aa (n7". .) I
Ez (17" I "n") [Rowland et al., 2018, Proposition 3]. Applying
LSÁ solving the equation yields tabular categorical TD learning, and the iteration rule is given by
(s() <- (s:) a(ne-1(8:) I (bre)) # nh- 11-16811)]
ne(8) < Th- T)-
aa + 8t-
Linear-Categorical TD Learning
In this section, we propose Linear-CTD algorithm (Eqn. (13)) by combining the linear function approximation (Eqn. (3)) with the categorical parametrization (Eqn. (7)). We first introduce the space of linear-categorical parametrized signed measures with total mass 1:
goiin Φ. = { [o = (mo())) : no(s)) = Ek Pi: (s: S2n, (0(0) ....., ...,0((K1)))) E Bad
(010))....,0(
which is an affine subspace of ( (grinnss Here pa (8; [x; - (s; (); and
Fa(s: 0) (bs) ((k)
T(k)) (k 1)/(K + 1) for k E 10. K K--1 -}
(10)
CDF of (s) Iy (F: = 0, Fx s; = 1). where (x+)(++) is the CDF of the discrete uniform distribution // on {xo..... ,.o
Hx used for normalization In many cases, especially when formulating and implementing algorithms, it is much more convenient and efficient to work with the matrix version of the parameter o: (0(0).... Q(K-1) d rather than with @=vec (0).
We define the linear- catecrrcaeeooicall projection operator II k. (gasign) gggiinn as follows:
IT Ø == argmin
2 (1, 70), bn (g
no E spoo $.8
H is in fact orthogonal projection (see Proposition D.2), and thus is non-expansive. The following proposition characterizes 11TT
Φ whose proof can be found in Appendix D.2.
Proposition 3.1. For any 1 (gosign) 22, is uniquely given by ny. where θ = vec (),
θ Ô-Σ E.-n (K)+)))))) CCT, C= MHKEETK K (11)
Here Pr (1)) K-1 the vector that identifies IT((()) defined Eqn. (8).
The following proposition characterizes the categorical projected Bellman operator II when acts on the parametrized n7. whose proof can be found in Appendix D.3.
Proposition 3.2. For any 0 E BdK and E S, we abbreviate pT njo (s) as (s()), then
(s) (p: 0)) K-1 =E (0T = + K+1 Σ 1=
9 (r(o)) su
where for any r E [0, and j. E 10,
(r. = (1 [r yr. - -xx//K()4+ g,(r) (g), E ERK G(r) = (g..))..... (r)
9K-
Kxk, C-):=G(m - 1 @ gx (r) e Kx
'The normalizer v is indispensable for achieving tight sample complexity bound, and it also guarantees our estimator has total mass 1, making the estimator more interpretable.
5

---

Since IT T is coonttractoo in (gesi $. k2,p-) (T is conttactto [Bellemare et al., 2023, Lemma 9. 9.14J, in Theorem 4.1, using Proposition 3.1 and 3.2. we generalize the linear projected Bellman equation (Eqn. (4)) to the distributional setting. The proof can be found in Appendix D.4.
Theorem 3.1. The linear-categoricdd projected Bellman equation no= Φ. T no admits unique solution nle where the matrix parameter O* is the unique solution to the linear system for DeRdx
Mg - (bb-ae.._pb))((()))c))))))))0))0
1k)
(12)
In analogy to the approximation bounds of | VE- ll and WT. (n" .7 k, the following lemma answers how close Tle- to 17", whose proof can be found in Appendix D.5.
Proposition 3.3 (Approximation Error of 17-)). It holds that
W (n² 170- - K 1 3 + (1 (1
" I²,
where the first error term error term (1- 2 # (II kn II x²) is due to the additional linear function approximation.
K- 7)
3 is due the calegorical parametrization, and the second
As before, we solve Eqn. (12) by LSA and get Linear CTD given the streaming data {(5t, Qe, Te)
Linear-CTD:
0 0 - (((() ( - ( 0 (L-(CG())))
(13)
-(K+1) (E (r:) (nt)
for any 1, where is the constant step size. is easy to verify that, in the special case of the tabular setting (ø(s) (1(s= Linear-CTD is equivalent to tabular cate- gorical TD learning (Eqn. (9)). In this paper, we consider the Polyak-Rupper tail averaging 0r == -12+1))--E-- Og (we use an even number 7) as in the analysis of Linear-TD in Samsonov et al. [2024b]. Standard theory of LSA [Mou et al., 2020] says under some conditions, we take an appropriate step size a, Or will converge to the solution with rate T as T + oo.
In Figure 1, we empirically validate the convergence of Linear-CTD through numerical experiments. See Appendix for details of the numerical experiments.
9
8
7
==
1og
6
K=
K=5
5
K
4
500
0
10000
20000
30000
40000
Iterations
Figure 1. Convergence results under varying K for our Linear-CTD algorithm with step size a = 0.01. These curves exhibit similar trends, demonstrating our algorithm's 's robustness across different K values.
Remark (Comparison with Existing Linear Distributional TD Learning Algorithms). Our Limear-CTD can be regarded as preconditioned version of vanilla stochastic semi semi- gradient descent (SSGD) with the probability mass function (PMF) representation [Bellemare al., 2023, Section 9.6]. See Appendix D.6 for the PMF representation, and Appendix D.7 for a self-contained derivation of SSGD with PMF representations. The preconditioning technique a commonly used methodology accelerate solving optimization problems by reducing the condition number. We precondition the

---

vanilla algorithm by removing the matrix C C (see Eqn. (26)). whose condition number scales with K² (Lemma H.2). By introducing the preconditioner C)- our Linear-CTD (Eqn. (13)) can achieve convergence rate independent of K, which the vanilla form cannot achieve. See Remark 5 and Appendix G for theoretical and experimental evidence respectively.
We comment that Linear-CTD (Eqn. (13)) is equivalent to SSGD with CDF representation, which was also considered in Lyle et al. [2019]. The difference is that our Linear-CTD normalizes the distribution so that the total mass of return distributions always be i, while the algorithm in Lyle et al. [2019] does not, See Appendix D.7 for self-contained derivation of SSGD with CDF representations.
The previously mentioned works, as well as our Linear-CTi CTD, are all limited by the use of signed mea- sures (Eqn. (7)). which makes them less interpretable. Bellemare et al. [2023, Section 9.5] proposed a softmax-base basaa iinear-categoriall algorithm, which is more closer to the practical C51 algorithm [Bellemare al., 2017] and it uses standard probability measures. However, the nonlinearity due to softmax makes it difficult for analysis. We will leave the analysis for it as future work.
Remark 2 (Linear-CT) is mean-preserving)) A key property of Linear-CTD is mean preservation. That is, if we identical initializations (BGGTT, [G] = Vv and an identical data stream to update in both Linear-CTD and Linear-TD, it follows that Eg-no [G]
V for all However, the mean-preserving property does not hold for the SSGD with the PMF representation. In this sense, our Linear-CTD indeed the generalization of Linear-TD. See Appendix D.8 for details.
4 Non-Asymptotic Statistical Analysis
In our task, the quality of estimator nØr measured by We--htttted 1 1-Wasserstein error Wi, n²). By triangle inequality, the error can be decomposed into the approximation error and the estimation error: W Wi,,, (n7": N (n7" no-) Mo-)+Wi.ppp (No. Nøz): Proposition 3.3 already provided an upper bound for the approximation error W Wi,a, (17". 7o- so suffices control
the estimation error W1, (me- ), denoted ((r)).
In the following theorem, we give non-asymptotic convergence rates of [(Or). We start from the gen- erative model setting, i.e., in the 1-th iteration, collect samples st (r), ae (-1ss), (re, s() ~ Pl,,, az) from the generative model, and we replace 81+1 with s/ in Eqn. (13). We give L and begh-probbbility convergence results in this setting. These results can be extended the Markovian setting, i.e., using the streaming data {(s, at, rr) rt)]==0
4.1 1 Convergence
We first provide asy-ssymio convergence rates of (C(O) which do not grow with the number of supports K. The L (p > 2) convergence results can be found Theorem E.1.
Theorem 4.1 (T.2 Convergence). For any 2 (1 -1 and a E (0, (1 - /7)/76), holds that
110"11,, 11 +1
110 l1 +1
E (V(((((()))) -
(1+ (1 M)min
VT(11 7) Aminn
T Tva(1 7) Amin
(1 fa((1 Amin
T/2
110- 0 11
T Taa(1 7) Ammi
Va
(1177) Amin
where 10*1,,, = (K(1->) 100 IIx and 110 0 001lvs = vK(1-7) 100 1.
In the upper bound, the first term of order 112 is dominant. The second term of order - has a worse dependence on -1 leading a sample size barrier (Eqn. (15)). The third term, corresponding to the initialization error, decays at a geometric rate and can be thus ignored. To prove Theorem 4.1, we conduct fine- -eeigrr analysis of the linear- categoriia Bellman equation and apply the exponential stability argument [Samsonov et al., 2024b]. We outline the proof in Section 5. In Remark 3, we compare our Theorem 4. I with the L22 convergence rate of classic Linear-TD and conclude that learning the distribution of the return is as easy as learning its expectation (value function) with linear function approximation. Rowland et al. [2024b], Peng et al. [2024) discovered this phenomenon in the tabular setting, and we extended it to the function approximation setting.
7

---

Remark (Comparison with Convergence Rate of Linear-TD). The only difference between our Theorem 4.1 and the tight L convergence rate of classic Linear-ID (see Appendix C.2) lies in replacing ||a/" I| (resp. || in Linear-TD with ||0 (resp. ||0g- 1100-0'1lv,) 11,,)
We claim that
the two pairs should be of the same order, respectively. Note that 100 and |10--0 ||,, are of order (((--7) if no- (s) and noo (s) are valid probability distributions for all E S. This is because in
E [) 0 E 10- 0o)- While in Linear-TD. this case, Fs:s;; 0) ((s) [k()))) (k+t))(K
|| and all E S and ap E {,b", 3h0)- It is thus reasonable to consider the two pairs with the same order, respectively. Similar arguments also hold other convergence results presented this paper. Therefore, in this sense, our results match those of Linear-TD.
are also of order (((-))=)) if (s(
Ø(s) 2/ (1-7) for
One can translate Theorem 4.1 into a sample complexity bound. Corollary 4.1. Under the same conditions as Theorem 4. I, for any e > O, suppose
10 11 +1
(1+ (1
10-1lv; +1
T
=²(1 2m2mm
Mmmm
Eva(l 7) Ammin
|00 0 flll
(Va
1 Va V(1-2)amnn
+ a(1 Nuiin
log
+ +log
TVa(L-n)V\nin
E
Then holds that (3/2((C((0T))] S
If we take the largest possible instance-indepeddeer step size, i.e. a (1 (1- and consider (0, 1), we obtain the sample complexity bound
Step Size.
T T= 5 2 2 7) (10*11 + +1))
(14)
Optimal Instance-eppende Step Size. If take the optimal instance-dependee step size 2 (1 - 7))mmn which involves the unknown Amin, we obtain better sample complexity bound
T T= δ ((=-2 ( (1- (1-7)-2>> (10-11 + +1))
(15)
There is sample size barrier in the bound, that is, the dependence on Amiu the optimal Xmi only when == (vmin)) or equivalently, we require a large enough (independent of E) update steps T.
These results match the recent results for classic Linear-TD with a constant step size [Li et al., 2024b, Samsonov et al., 2024b]. It is possible to break the sample size barrier Eqn. (15) as in Linear-TD by applying variance-reduction techniques [Li et al., 2023a]. We leave it for future work.
4.2 Convergence with High Probability and Markovian Samples
Applying the IP convergence result (Theorem E. with p = log and Markov's inequality, we immediately obtain the high-probability convergence result.
Theorem 4.2 (Hiigh-rooabblity Convergence). For any > 0 and å E (0, 1), suppose K 2 (1-7)-1 (0, (1 vT) log (o'((T/8")) and
110 110-11 +1
(1+ (1+ (1+(17))mmm))
a log coogg
log 1 J eva(1 7)3 Amin
1 111 +1 log a(1 111-))mmm Ami
log
1100 00
T T= δ
e2 - 7mmmm
Then with probability at least - 8, it holds that C (Ør) < Here, the Ō.) does not hide polynomials of 1og(1/8) (but hides logarithm terms of (1/8)).
Again, we will obtain concrete sample complexity bounds as Eqn. (14) or Eqn. (15) if we use different step sizes. Compared with the theoretical results for classic Linear-TD, results match Samsonov et al. [2024b, Theorem 4]. Samsonov et al. [2024b] also considered the constant step size, but obtained worse dependence on log (1/8) than Wu et al. [2024, Theorem 4] which uses the polynomial-decaying step size a₃ a with 3 (1/2, instead.

---

Remark 4 (Markovian Setting). Using the same argument as in the proof of Samsonov et al. [2024b. Theorem 6], one can immediately derive a thigh-proborn probabiiy sample complexity bound in the Markovian setting. Compared with the bound in the generative model setting (Theorem 4.2), the bound in the Markovian setting will have an additional dependency on mmix (og(T/B) /6), where bmix is the mixing time of the Markov chain {8t t= in S. We omit this result for brevity.
5 Proof Outlines
In this section, we outline the proofs of our main results (Theorem 4.1). We first state the theoretical properties of the linear-categorical Bellman equation and the exponential stability of Linear-CTD. Finally, we highlight some key steps in proving these results.
5.1 Vectorization of Linear-CTD
In our analysis, it will be more convenient to work with the vectorization version of the updating scheme of Linear-CTD (Eqn. (13)):
0, P.400-1-00 (A,0011- A,= [ (p())))(s)) -(CG)))))))))(((())))))))))))))) (d(o))))))
"),,
b. (K + (C (r:)
1x)] s (ss)).
(16)
We denote by A and b the expectations of A, and b, respectively. Using exponential stability arguments, we can derive an upper bound for | - 0l) [|- The following lemma further translates it an upper bound for [(OT) Wr. (No- No,). whose proof is given Appendix E.
Lemma 5.1. For any 0 E R holds that L(0) 2K 1/2 2
- A (0 0*)
min
5.2 Exponential Stability Analysis
First, introduce notations. Letting e, A denote by CA (resp. C.) the almost sure upper bound for max [| A+ [| (resp. ||e:l). and Σ. == E ee the covariance matrix of et. The following lemma provides useful upper bounds, whose proof is given in Appendix E.2.
Lemma 5.2. For any K 2 (1 7)-1, 11, it holds that
CAS4 S
C. S 4(110'1| + VK - 7)), tr tr.) S ( IXOE, + K[11 ((1-7)3).
Let r == -T aA)) for any a > and f N. The exponential stability Linear-CTD summarized in the following lemma, whose proof can be found in Appendix E.3.
Lemma 5.3. For any p 2 2, let a = (1 -Tmmmu//? and appoc = (1 /3/p1)((p)pp)))) (ap.ocP S 1/2). Then for any a E (0, ap,cc). u E RaK and t E N, it holds that E | u u||l]]
= (1 aa)² ||ul|.
The following theorem states the I convergence of || A 0') based on exponential stability arguments. For the general L convergence, please refer to Samsonov al. [2024b, Theorem 2]. Theorem 5.1. (Samsonov al., 2024h, Theorem 1] For any (0, 022 it holds that
E -/01 tr Ee) VT
(1+ Cava 1+
/tr
T/2 (1 (1-aa)T/2
(1+ CA Vaa
+
1100-0%
Va
Tvaa
T
a
Combining these lemmas with Theorem 5.1, we can immediately obtain Theorem 4.1.
Remark 5 (Convergence of SSGD with PMF representation). In Appendix E.5, we give counterparts of these lemmas and Theorem 4.1 for vanilla SSGD with PMF representation. The results imply that the step size in the algorithm should scale with (1 (1- /)² and the sample complexity grows with K. In Appendix G.2, we verify through numerical experiments that as K increases, to ensure convergence, the step size of the vanilla algorithm indeed needs to decay at rate of K-2, In contrast, the step size of our Linear-CTD does not need to be adjusted when K increases. Moreover, we sind that L Linear-CTD empirically consistently outperforms the vanilla algorithm under different K.

---

5.3 Key Steps in the Proofs
Here we highlight some key steps in proving the above theoretical results.
Bounding the Spectral Norm of Expectation of Kronecker Products. In proving that the (())
can be upper-bounded by || | 0*)ll (Lemma 5.1), as well as in verifying the exponential stability condition (Lemma 5.3), one of the most critical steps is to show
Es.ne' (CG())))) & 2 1 (ss))(s))) "T
)]]
S VT, Vr E [0, 1].
(17)
By Lemma H.3, we have that ||Es,s' (2 1/2 1/2 ((s)((s') 2
for any r [0, 1]. In addition, one can check J One may speculate that the property | [ B2ll (Lemma A.3) is enough get the desired conclusion. However, the two ma- trices in the Kronecker product are not independent, preventing us from using this simple property to derive the conclusion. On the other hand, since we only have the upper bound Bg,e' ||
(t==011(11(0))000 ( 1/2 [],,, simply moving the expectation in Eqn. (17) outside the norm will lead to a loose d bound. To resolve this problem, we leverage the fact that the second matrix rank-1 and prove the following result. The proof can be found in the derivation following Eqn. (29).
Lemma 5.4. For any random matrix Y and random vectors æ. 2, suppose || S Cy almost surely, CdId, and E 1zz" 1 C=Idz for some constants Cr, Ca. C2 > 0. Then it holds that
E xx
E Y x T)[l SCV
Remark (Matrix Representation of Categorical Projected Bellman operator). The matrix CG(r)C also appears in Rowland al. [2024b, Proposition B.2] as the matrix representa- tion of the categorical projected Bellman operator IIxT" of a specific one-state MDP. As result, [ICG())C-1| because Ig is a /-cootrraction in (9,b)) Our Lemma H.3 provides new analysis by directly analyzing the matrix.
Bounding the Norm of b, In proving Lemma 5.2, the most involved step is to upper-bound | (Eqn. (16)). To this end, we need to upper-bound the following term:
K+1 - 1411 1x)ll. Vr E 1].
1
(18)
Term (18) is also related to the categorical projected Bellman operator TT. Specifically, let K+I Ex=o oz, be the discrete uniform distribution. One can show that Term (18) equals
[CC(P())))))) Pu)ll = 'K" 1/282 E2 (bra)))) S 'x" -1/2 ((bra) # (V), s 3v K(1 7),
where we used the fact that I is non-expansive and an upper bound for &a ((b,,) ) (v), v) when K (1 7) (Lemma H.4). The full proof can be found in the derivation following Equ. (30).
Conclusions
In this paper, we have bridged critical theoretical gap in distributional reinforcement learning by establishing the non-asymptotic sample complexity of distributional TD learning with linear function approximation. Specifically, we have proposed Linear-CTD, which is derived by solving the linear-categorical projected Bellman equation. By carefully analyzing the Bellman equation and using the exponential stability arguments, we have shown tight sample complexity bounds for the proposed algorithm. Our finite-sample rates match the state-of-ther sample complexity bounds for conventional TD learning. These theoretical findings demonstrate that learning the full return distribution under linear function approximation can be statistically as easy as conventional TD learning for value function estimation. Our numerical experiments have provided empirical validation of our theoretical results. Finally, we have noted that it would be possible to improve the convergence rates by applying variatce-reducciion techniques or using polynomial-decayinn step sizes, which we leave for future work.
10

---

Acknowledgments and Disclosure of Funding
The authors would like to thank the reviewers for their constructive feedback, which helped us improve the quality of our work.
This work has been supported by the National Key Research and Development Project of China (No. 2022YFA1004002 and the National Natural Science Foundation of China (No. 12271011 and No. 12350001).
References
B. Ávila Pires, M. Rowland, D. Borsa, Z. D. Guo, K. Khetarpal, A. Barreto, D. Abel, R. Munos, and W. Dabney. Optimizing return distributions with distributional dynamic programming. Journal of Machine Learning Research, 26 (6988))1-990, 2025. URL http:://jmml http:///mn...or/ppesss////
255-210.2 html.
N. Bäuerle and Ott. Markov decision processes with average-value-at-e criteria. Mathematical Methods of Operations Research, 74:361-379, 2011.
M. G. Bellemare, W. Dabney, and R. Munos. A distributional perspective reinforcement learning. In International conference on machine learning, pages 449-458. PMLR, 2017.
M. G. Bellemare, N. Le Roux, P. S. Castro, and S. Moitra. Distributional reinforcement learning with linear function approximation. In The 22nd International Conference Artificial Intelligence and Statistics, pages 2203-2211. PMLR, 2019.
M. G. Bellemare, W. Dabney, and M. Rowland. Distributional Reinforcement Learning. MIT Press, 2023. http://www.discccbbttiionnall distrrbbutionallll..cco
D. Bertsekas and N. Tsitsiklis. Neuro-dynamic programming: an overview. In Proceedings of 1995 34th IEEE conference on decision and control, volume 1, pages 560-564. IEEE, 1995.
J. Bhandari, D. Russo, and R. Singal. A finite time analysis of temporal difference learning with linear function approximation. In Conference on learning theory, pages 1691-1692. PMLR, 2018.
M. Böck and C. Heitzinger. Speedy categorical distributional reinforcement learning and complexity analysis. SIAM Journal on Mathematics of Data Science, 4(2):675-693, 2022. doi: 10.1 1377/ 20M1364436. URL httpt: wrg/10gg//ooooo 10/g/10..10 1137 10111117711111 2201///0003344436.
K. Chen. Matrix preconditioning techniques and applications. Cambridge University Press, 2005.
Z. Chen, S. T. Maguluri, S. Shakkottai, and K. Shanmugam. A Iyapunov theory for samite-sample guarantees of markovian stochastic approximation. Operations Research, 13(2):33521136) 2024.
T. Cho, S. Han, K. Lee, S. Ju, D. Kim, and J. Lee. Tractable and provably efficient distributional reinforcement learning with general value function approximation. arXiv e-prints, pages arXiv- 2407, 2024.
Y. Chow and M. Ghavamzadeh. Algorithms for evar optimization in mdps. Advances in neural information processing systems, 27, 2014.
W. Dabney, G. Ostrovski, D. Silver, and R. Munos. Implicit quantile networks for distributional reinforcement learning. In International conference machine learning, pages 1096-1105. PMLR, 2018a
W. Dabney, M. Rowland, M. Bellemare, and R. Munos. Distributional reinforcement learning with quantile regression. Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018b.
G. Dalal, B. Szörényi, G. Thoppe, and S. Mannor. Finite sample analyses for td (0) with function approximation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018
Y. Duan and M. Wainwright. finite-sample analysis of multi-step temporal difference estimates. In Learning for Dynamics and Control Conference, pages 612-624. PMLR, 2023.

---

A. Durmus, E. Moulines, A. Naumov, and S. Samsonov. Finite-timg booh-probability bounds for polper-ruppert ruppert averaged iterates of linear stochastic approximation. Mathematics of Operations Research, 2024.
D. Freirich, T. Shimkin, R. Meir, and A. Tamar. Distributional multivariate policy evaluation and exploration with the bellman gan. In International Conference on Machine Learning, pages 1983-1992. PMLR, 2019.
C. D. Godsil. Inverses of trees. Combinatorica, 5:33-39, 1985.
R. A. Horn and C. R. Johnson. Topics in matrix analysis. Cambridge university press, 1994.
D. Huo, Y. Chen, and Q. Xie. Bias and extrapolation in markovian linear stochastic approximation with constant stepsizes. In Abstract Proceedings of the 2023 ACM SIGMETRICS International Conference on Measurement and Modeling of Computer Systems, pages 81-82, 2023.
T. Kastner, M. A. Erdogdu, and A. Farahmand. Distributional model equivalence for risk-sensitive reinforcement learning. Advances in Neural Information Processing Systems, 36:56531-56555 2023.
T. Kastner, M. Rowland, Y. Tang, M. A. Erdogdu, and A. massoud Farahmand. Categori- cal distributional reinforcement learning with kullback-leibler divergence: Convergence and asymptotics. In Forty-second International Conference on Machine Learning, 2025. URL https://openrerier
C. Lakshminarayanan and C. Szepesvari. Linear stochastic approximation: How far does constant step-size and iterate averaging go? In International conference on artificial intelligence and statistics, pages 1347-1355. PMLR, 2018.
G. Li, C. Cai, Y. Chen, Y. Wei, and Y. Chi. Is q-leanning minimax optimal? tight sample complexity analysis. Operations Research, 72(1):222-236, 2024a.
G. Li, W. Wu, Y. Chi, C. Ma, A. Rinaldo, and Y. Wei. High-probability sample complexities for policy evaluation with linear function approximation. IEEE Transactions on Information Theory, 2024b.
Li, G. Lan, and A. Pananjady. Accelerated and instance-optimal policy evaluation with linear function approximation. SIAM Journal Mathematics of Data Science, 5(1):174-200, 2023a.
X. Li J. Liang, and Z. Zhang. Online statistical inference for nonlinear stochastic approximation with markovian data. arXiv preprint arXiv:2302.0766 2023b.
X.L.. Li. Preconditioned stochastic gradient descent. IEEE transactions neural networks and learning systems, 29(5):1454- 1466, 2017.
S. H. Lim and I. Malik. Distributional reinforcement learning for risk-sensitive policies. Advances in Neural Information Processing Systems, 35:30977-30989, 2022.
Y. Luo, G. Liu, H. Duan, O. Schulte, and P. Poupart. Distributional reinforcement learning with monotonic splines. In International Conference on Learning Representations, 2022. URL https: //openrvieww www.:nee/oorrree dd
C. Lyle, M. G. Bellemare, and S. Castro. comparative analysis of expected and distributional re- inforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 4504-4511, 2019.
A. Marthe, A. Garivier, and C. Vernade. Beyond average return in markov decision processes. In sevethh Conference on Neural Information Processing Systems, 2023. URL https: // //opnreview
M. Moghimi and H. Ku. Beyond CVar: Leveraging static spectral risk measures for enhanced decision- making distributional reinforcement learning. Forty-second International Conference on Machine Learning, 2025. URL http:p::/oopeeeeeeeeeeeeerrrrrrr
12

---

T. Morimura, M. Sugiyama. H. Kashima, H. Hachiya, and T. Tanaka. Nonparametric return dis- tribution approximation for reinforcement learning. In Proceedings of the 27th International Conference on Machine Learning (ICML-10), pages 799-806, 2010.
Mou, J. Li, J. Wainwright, P. L. Bartlett, and M. Jordan. On linear stochastic approximation: Fine-grained polyak-ruppert and non-asymptotic concentration. In Conference on Learning Theory, pages 2947-2997. PMLR, 2020.
W. Mou, A. Pananjady, M. Wainwright, and P. Bartlett. Optimal and instance-dependent guarantees for markovian linear stochastic approximation. In Conference on Learning Theory, pages 2060-2061. PMLR, 2022.
T. Nguyen-Tame S. Gupta, and S. Venkatesh. Distributional reinforcement learning via moment matching. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9144-9152, 2021.
E. Noorani, C. N. Mavridis, and J. S. Baras. Exponential td learning: A risk-sensitive actor-critic reinforcement learning algorithm. 2023 American Control Conference (ACC), pages 4104-4109. IEEE, 2023.
G. Patil, L. Prashanth, Nagaraj, and D. Precup. Finite time analysis of temporal difference learning with linear function approximation: Tail averaging and regularisation. In International Conference on Artificial Intelligence and Statistics, pages 1448. PMLR, 2023.
Y. Peng. L. Zhang, and Z. Zhang. Statistical efficiency of distributional temporal difference learning. In The Thirty-eighth Annual Conference Neural Information Processing Systems, 2024. URL http://oruumnideeeutsSntYYE
B. T. Polyak and A. B. Juditsky. Acceleration of stochastic approximation by averaging, SIAM journal on control and optimization, 30(4):838-855, 1992.
Z. Qi. C. Bai, Z. Wang, and L. Wang. Distributional off-policy evaluation in reinforcement learning. Journal of the American Statistical Association, jun 2025. doi: 10. 10.1080001622455.222525556696699 Forthcoming.
C. Qu, S. Mannor, and H. Xu. Nonlinear distributional gradient temporal-differenn learning. In K. Chaudhuri and R. Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 5251-5260 PMLR, 09 15 Jun 2019. URL htthttp:: http///pooeeeensee//eee mlr. www.ppeessvvvee patpss//v////uuluu /qul///uibbbbbb/.. html.
M. Rowland, M. Bellemare, W. Dabney, R. Munos, and Y. W. Tch. An analysis of categorical distributional reinforcement learning. In International Conference Artificial Intelligence and Statistics, pages 29-37. PMLR, 2018.
M. Rowland, R. Dadashi, S. Kumar, R. Munos, M. G. Bellemare, and W. Dabney. Statistics and samples in distributional reinforcement learning. International Conference on Machine Learning, pages 1522-55384 PMLR, 2019.
M. Rowland, Y. Tang, C. Lyle, R. Munos, M. G. Bellemare, and W. Dabney. The statistical benefits of quantile temporal-difference learning for value estimation. In International Conference on Machine Learning, pages 29210-29231. PMLR, 2023.
M. Rowland, R. Munos, M. G. Azar, Tang, G. Ostrovski, A. Harutyunyan, K. Tuyls, M. G. Bellemare, and W. Dabney. An analysis of quantile temporal-different learning. Journal of Machine Learning Research, 25:1-47, 2024a.
M. Rowland, L. K. Wenliang, R. Munos, C. Lyle, Y. Tang, and W. Dabney. Near-minimax-optimall distributional reinforcement learning with generative model. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024b. URL https://openreviae net/forum?id:JJKBG 1d4ib.
D. Ruppert. Efficient estimations from a slowly convergent robbins-monro process. Technical report, Comell University Operations Research and Industrial Engincering, 1988.
13

---

S. Samsonov. E. Moulines, Q.-M. Shao, Z.-S. Zhang, and A. Naumov. Gaussian approximation and multiplier bootstrap for pupyak-ruppert averaged linear stochastic approximation with applications to TD learning. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024a. URL htttt::://pppeeeeeeeeeeee
JAs
S. Samsonov, D. Tiapkin, A. Naumov, and E. Moulines. Improved high-probability bounds for the temporal difference learning algorithm via exponential stability. In The Thirty Seventh Annual Conference on Learning Theory, pages 451 1-4547. PMLR, 2024b.
D. Serre. Matrices: Theory and Applicattion. Graduate texts in mathematics. Springer, 2002. ISBN 9780387954608. URL attpps://booki google. com co..hh//ooossrdddddiiiYYY FURRRUUUKGG
G. Shen, R. Dai, G. Wu, S. Luo, C. Shi, and H. Zhu. Deep distributional learning with non-crossing quantile network. arXiv preprint vrXiv:2504.08215 2025.
R. Srikant and L. Ying. Finite-time error bounds for linear stochastic approximation andtd learning. In Conference Learning Theory, pages 2803-2883 PMLR, 2019.
K. Sun, Y. Zhao, W. Liu, B. Jiang, and L. Kong. Distributional reinforcement learning with regularized wasserstein loss. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024. URL https: / http//ppeeeeveeeeeeeeeeeeeeeeeee net http://eeeeee.eeeeeeeeeeeeeeeeee foun:
R. S. Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3:9 44, 1988.
Y. Tang, R. Munos, M. Rowland, B. Avila Pires, W. Dabney, and M. Bellemare. The nature of temporal difference errors in multi-step distributional reinforcement learning. Advances in Neural 2022.
Information Processing Systems, 35:30265-3000
Y. Tang, M. Rowland, R. Munos, B. Á. Pires, and W. Dabney, Off-policy distributional q (lambda): Distributional rl without importance sampling. arXiv preprint arXiv:2402.05766. 2024.
J. Tsitsiklis and B. Van Roy. Analysis of Advances in neural information processing systems, 9, 1996.
difffferefce learning with function approximation.
R. Vershynin. High- Ligh--Dimessooal Probability: An Introduction with Applications in Data Science. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018. doi: 109233559...
H. Wang and X. Y. Zhou. Continuous-time mean-variance portfolio selection: A reinforcement learning framework. Mathematical Finance, 1074) 2020.
K. Wang, K. Zhou, R. Wu, N. Kallus, and W. Sun. The benefits of being dis- tributional: Small-loss bounds for reinforcement learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neu- ral Information Processing Systems, volume 36, pages 2275-2312. Curran Associates. Inc., 2023. URL https::/proceetee 06f 06ff5c2221 htt.aae22ee66011100227777778668 Papeaa Conference, pdf.
neurips. cc/paper_rppppppeepeeeeeeeee fotrr_fill//eeleeefr halles/paper/022/1112
K. Wang. O. Oertell, A. Agarwal, N. Kallus, and W. Sun. More benefits of being distributional: Second-order bounds for reinforcement learning. In Forty-first International Conference Machine Learning, 2024. URL
L. Wang, Y. Zhou, R. Song, and B. Sherwood. Quantile-optimal treatment regimes. Journal of the American Statistical Association, 1133523).1243-11252 2018.
L. K. Wenliang, G. Deletang, M. Aitchison, M. Hutter, A. Ruoss, A. Gretton, and M. Rowland. Distributional bellman operators mean embeddings. In Forty-first International Conference on Machine Learning, 2024. URI hwl.lttpp::/ppeeeeeievv.. 3
H. Wiltzer, J. Farebrother, A. Gretton, and M. Rowland. Foundations of multivariate distributional reinforcement learning. The Tiirty-eightt Annual Conference on Neural Information Processing Systems, 2024a. URL https://opperevies. httpr//www....eefooorrrrrraaaaa.
14

---

H. Wiltzer, J. Farebrother, A. Gretton, Y. Tang, A. Barreto, W. Dabney, M. G. Bellemare, and M. Rowland. A distributional analogue the successor represtttion...... International Conference on Machine Learning (ICMI., 2024b.
R. Wu, M. Uehara, and W. Sun. Distributional offline policy evaluation with predictive error guar- antees. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, Proceedings of the 40th International Conference Machine Learning. volume 202 of Pro- ceedings of Machine Learning Research, pages 37685-37712. PMLR, 23-29 Jul 2023. URL https: sttaps:/pproceedings mlr
wwwp://ww33s..mmmm html.
W. Wu, Li, Y. Wei, and A. Rinaldo. Statistical Inference for Temporal Difference Learning with Linear Function Approximation. arXiv preprint arXiv:2410. 16106, 2024.
D. Yang, L. Zhao, Z. Lin, T. Qin, J. Bian, and T.-Y. Liu. Fully parameterized quantile function for distributional reinforcement learning. Advances neural information processing systems, 32, 2019.
Y. Yue, Z. Wang, and M. Zhou. Implicit distributional reinforcement learning. Advances in Neural Information Processing Systems, 33:7135-7147, 2020.
H. Zhang and Ding. On the kronecker products and their applications. Journal of Applied Mathematics, 2013(1):296185, 2013.
L. Zhang, Y. Peng, Liang, W. Yang, and Z. Zhang. Estimation and Inference in Distributional Reinforcement Learning. The Annals of Statistics, 2025.
F. Zhou, Z. Zhu, Q. Kuang, and L. Zhang. Non-decreasing quantile function network with efficient exploration for distributional reinforcement learning. In Z.-H. Zhou, editor, Proceedings of the Thirtieth International Joint Conference Artificial Intelligence, IJCAI-21, pages 3455-3461. International Joint Conferences on Artificial Intelligence Organization, 2021. doi: 10.24963/ 2jcai.2021/476. URL https: /ttp:///ooo///...////// 24963/ /aj0////jaai../ 2021/476.
jjcai.2021/476 Main Track.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We justify our claims in the abstract and introduction using rigorous proof.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations and future work Section Conclusions.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
16

---

Justification:
We provide full assumptions and proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We disclose all details of experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
17

---

Answer: [Yes]
Justification:
We provide the code in supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
connoollcyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips. wttpcccccppppuuiccceeeecceocccca sute iettssummssiioppooooooooeoo
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We disclose all details of experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: N/A.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
18

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide full information the computer resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttp///eerrpp.///ee cc htcp/pbiic/thhcccuudeeee.
Answer: [Yes]
Justification: Our research confirms with the NeurIPS Code of Ethics in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper only focuses the theory of RI., and there is societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
19

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper only focuses theory of RL, there is such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We properly use and cite these assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, haperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
20

---

Answer:
[NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/codd/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
21

---

The core method development in our paper does not involve LLMs as any important, original, or non-standard components.
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non- non-standard components.
Please refer to our LLM policy (https: /ttpas////eeus./eeeeee https//c//ooooocnnaeeeeeece ( for what should or should not be described.
22

---

A Kronecker Product
In this section, we will introduce some properties of Kronecker product used in our paper. See Zhang and Ding [2013] for detailed treatment of Kronecker product.
For any matrices A E m=** and B IRxx, the Kronecker product defined as
B is matrix in gmpxnng
B
012 B
B
021
022 B
d2n
A B
0m1 B
02 B
Gyn B
Lemma A.1. The Kronecker product is bilinear and associative. Furthermore, for any matrices
B1, B2, B3, B4 such that B, B,Ba can be defined, it holds that (B1 @ B2) (B, @ B,) = (B, & (B2Ba) mmixed-produce property).
Proof. See Basic properties and Theorem 3 in Zhang and Ding [2013].
Lemma A.2. For any matrices B1, B2, B3 such that B, B, can be defined, it holds that vec B1 B2Bs) (Bg @ B1) vec (B2).
Proof. See Lemma 4.3.1 in Horn and Johnson [1994].
Lemma A.3. For any matrices B and Bz, holds that B1 & B₂l [ B2], (BO B₂) B @ BJ. Furthermore, B, and B2
Bi @ B2 also hnertheeorhgguuullgggggysveneeee
and [ @ B2) B2)- B, B
Proof. See Basic properties, Theorem 5 and Theorem in Zhang and Ding [2013].
Lemma A.4. For any K,d N and PSD) matrices B1, B3 R B2, BA E Radd with B, < B3 and B2 & Ba, it holds that B @ B2. B3 % B4 are also PSD matrices, furthermore, B x B2 B Ba.
Proof. Consider the spectral decomposition B; Q.D.Q. for any i e [4], by Lemma A.1 and Lemma A.3, we have
B & B2) (Q1 Q2) (Di > D2) (Qi Q2)
and
(Ba & B.) = (Q3 & Q+) DD3 3 D.) Q23 Q.)
also spectral decomposition of [B] @ B2) and [B3 B₁) respectively. is casy to see that they are PSD. Furthermore,
(Ba B4) (B, B2) [B3 @ B,) (B3 @ B₂) [B3 B2) (Bi Bz)]
Bg @ (Ba B2) (B - B 0 Bz)
0.
Lemma A.5. For any K, d, d2, d2 E N, vectors u, E Rdd and matrices B E Kxd, B2 E Rdxa
By R Kx it holds that
(Ix 8 u) B B 0u,
B. (Ix & v) B2 T.
(IK @ u) B3 [I & T)T B3 (wv7).
Furthermore, for any matrix B4 E [Rd xd2 we have
B 2 B (B &
23

---

Proof. Let u = (u4) B1 - (biy)
== then
0g
0d
0d
Od
0a
Od
b11
bik
(IK & u) B1
0a
0g
0a
bki
bKK
0a
0g
0a
u
binnd
bixua
b
brKua
bk
b
buuu
bkiu
bKKU
Bi 69 u.
Hence
B2 [Ix 2 v)
(IK 2 v) @ B
[B 00 B2 820r 8T
And in the same way,
(Ix 2 u B3 [I 80)
BB3 8 8
B3 3 2 B3 (uv")
Furthermore,
(B 8
[Ix 2 0 = (Ix (B, = (.) @ u.
B Related Work
Comparisons of Theoretical Results with the Previous Work In Table 1, we summarize our main theoretical results and comparing them with prior work. the table, when the task is policy evaluation, the sample complexity is defined in terms of the μl we-weighted L norm the measure of error; when the task is distributional policy evaluation, the sample complexity is defined terms of the we--hhhtteed W1 metric as the measure of error. The table gives clear comparison of theoretical results with prior work.
Distributional Reinforcement Learning. Distributional TD learning was first proposed in Belle- mare al. [2017]. Following the distributional perspective in Bellemare et [2017], 71 Qu al. [2019] proposed distributional version of the gradient TD learning algorithm, Tang et al. [2022] proposed distributional version of multi-step TD learning. Tang et [2024] proposed distributional version of off-policy Q(A) and TD(A) algorithms, and Wu et al. [2023] proposed distributional version of fitted Q evaluation to solve the distributional offline policy evaluation problem. Qi ct al. [2025] also considered the distributional off-policy evaluation problem. Wiltzer et al. [2024b] proposed an approach for evaluating the return distributions for all policies simultancously when the reward is deterministic or in the inite-horizon setting. Wiltzer al. [2024a] studied distributional policy evaluation the multivariate reward setting and proposed corresponding TD learning algorithms. Be- yond the tabular setting, Bellemare et al. [2019], Lyle al. [2019], Bellemare al. [2023] proposed various distributional TD learning algorithms with linear function approximation under different parametrizations.
series of recent studies have focused on the theoretical properties of distributional TD learning. Rowland et al. [2018], Böck and Heitzinger [2022], Zhang et al. [2025], Rowland et al. [2024a.b],

---

Paper
Sample Complexity
Method
15 (1-2) Amin
+1
Samsonov et al. [2024b]
Linear-TD
Amnin
1|40
12gg +1
Li et al. [2023a]
Linear-TD with Variance Reduction
(1(1-)) Ammin
Bellemare et al. (2023, Section 9.7]
Contraction Analysis
Vanilla Linear-
This Work (Theorem E.2)
K K* 18, +1)
K² Amin
Vanilla Linear-CTD
(1-2)) Amin
go
100 1R (I-2)² Amin
This Work (Theorem 4.1)
Linear-CTD
Amntn
Table Sample complexity of algorithms for solving policy evaluation and distributional policy evaluation. Here, Vanilla Linear-CTD refers to the stochastic semi-gradient descent algorithm with the prohability mass function representational (see Eqn. (27) in Appendix D.7). The contraction analysis [Bellemare et al., 2023, Section 9.7] means that they provided contraction analysis for the dynamic programming version of Vanilla Linear-CTD algorithm.
Peng et al. [2024], Kastner al. [2025] analyzed the asymptotic and non-asymptotic convergence of variants) in the tabular setting. Among these works, distributional TD learning (or its model-base
Rowland et al. [2024b]. Peng et al. [2024] established that in the tabular setting. learning the full return distribution is statistically as easy as learning its expectation in the model-basee and model-free settings, respectively. And Bellemare et al. [2019] provided an asymptotic convergence result for categorical TD learning with linear function approximation.
Beyond the problem of distributional policy evaluation, Rowland et al. [2023], Wang et al. [2023, 2024] showed that theoretically the classic value-based reinforcement learning could benefit from distributional reinforcement learning. Bäuerle and Out [2011], Chow and Ghavamzadeh [2014], Marthe al. [2023], Noorani et al. [2023], Moghimi and Ku [2025], Ávila Pires et al. [2025] considered optimizing statistical functionals of the return, and proposed algorithms to solve this harder problem.
Stochastic Approximation. Our Linear-CTD falls into the category of LSA. The classic TD learning, one of the most classic LSA problems, has been extensively studied [Bertsekas and Tsitsiklis, 1995, Tsitsiklis and Van Roy, 1996, Bhandari et al., 2018, Dalal et al., 2018, Patil et al., 2023, Duan and Wainwright, 2023, Li et al., 2024a,b, Samsonov et al., 2024a, Wu et al., 2024]. Among these works, Li et al. [2024b], Samsonov et al. [2024b] provided the tightest bounds for Linear-TD with constant step sizes, which is also considered in our paper. While Wu et al. [2024] established the tightest bounds for Linear-TD with pplynomial-decayinn step sizes.
For general stochastic approximation problems, extensive works [Lakshminarayanar and Szepesvari, 2018, Srikant and Ying, 2019, Mou et al., 2020, 2022, Huo et al., 2023, Li et al., 2023b, Durmus et al., 2024, Samsonov et al., 2024b, Chen et al., 2024] have provided solid theoretical understandings.
C Omitted Results and Proofs in Section 2
C.1 Linear Projected Bellman Equation
It is worth noting that, I1; [R-S 11-111mm.) - (Yo: Il- p1-l)p.) is orthogonal projection.
We aim to derive Eqn. (4). It is easy to check that, for any V E IRS, ny, V is uniquely give by V. where
- (s)s
Hence, by the definition of Bellman operator (Eqn. (1)), 2b- is the unique solution to the following system of linear equations for ab E Idd
v/ M> Esnna (0(s) T + [s)]
Σ
3 E (0(s) EE) Iro so = s] + [T T 8o = s] $)
Σ x Eyys' ( t E., (p(s)))
25

---

or equivalently,
EE - -Bu, (((s))(s'))) xb Es, (d(s)r)
C.2 Convergence Results for Linear TD Learning
is worthy noting that, Linear-TD is equivalent to the stochastic semi-gres ammi-gradient descent (SSGD) update.
In Linear -Tn TD, our goal is to find a good estimator & such that V
Vi++
S Samsonov et al. [2024b] considered the Polyak-Ruppert tail averaging TT = (T/2 1/2+10 1-1) -) vh and showed that in the generative model setting with constant step size = (1 7))
2/ Ex
||65
+1
-= -1-71 Amin
De
e2
Amin
is sufficient to guarantee that V Yo VΦ- ||, S
They also provided sample complexity bounds
when taking the instance-indepeenent (i.e., not dependent on unknown quantity) step size, and Markovian setting.
C.3 Categorical Parametrization is an Isometry
Proposition C.1. The affine space (ptten .22) is isometric with (RK /ex llT c). in the sense
that, for any Up p/pa E gpignn it holds that (g(pp: LK || P2lleTo where C' is defined in Eqn. (11).
Proof.
(1-7)-1
SlVp::Ups!
(F (x) Fvrr (x)) dx
0
(Fum (Ix) - Fur, (xx))
=
K=U
||C (p1 22)l1
= |p P2 1
C.4 Categorical Projection Operator is Orthogonal Projection
Proposition C.2. [Bellemare et al., 2023, Lemma 9.17] For any v E Toign and up E Dign it holds
that
t (u, up) = (v, [Kv) + & (IKRu
C,5 Categorical Projected Bellman Operator
The following lemma characterizing x is useful for both practice and theoretical analysis. Proposition C.3. For any m E (gosign) and S, it holds that
PT
(s) (s) =E 9K (ro) + (G(ro) - 1 @ gr())) P (2)( 80 =
-E G(ro) (s)(s)
K + 1K 80
E
9,
(ro)
80
K+1
26

---

And in the same way. for any r E [0, 1] and s' E S, it holds that
(ll')) = G(r) Pr((s)) K+1
1
1x
K K+ Σ g,(r),
where G and g is defined in Theorem 3.1.
This proposition is a special case of Proposition 3.2, whose proof can be found in Appendix D.3.
D Omitted Results and Proofs in Section 3
D.1 Linear-Categorical Parametrization an Isometry
Proposition D.1. The affine space
(geien $.K: f2, is isometric with (RAK .VLK 11 x in the
sense that, for any ne. n7. E S ø.K. holds that E2 (17. :T.) LK |10 - 02 1 80 Es.
Proof. For any De E grign
Φ we denote Fo(s) = (Fr(s;; 0)) E ERK then it holds that
1g (no, no,)
=1 Egnuy
Fo, (s) Fo, (s)
153 (+) -2) ( - -92) 2/}
=Ik tr
=Ik 10
0211 Tx00
D.2 Linear-Categorical Projection Operator
Proposition 3.1 an immediate corollary of the following lemma. For any v E PK define F, (F&s (v)) K-1 Pllnn(s) and F (s) FIxN(())-
K= (v xx)) K-1 R and for any n (gosien) s define pn(s))
Lemma D.1. For any 17 (gosin) 0 Rd and s S. it holds that
Vol (s)) (s)) 2(x (F) Fs(s))
Φ(s) O K-1 1x (K--
pm(s)) a
(19)
= 2(x)(())
Furthermore, it holds that
Volu 2n (No, 17) Egnpe Vo (n) [(}))
K 1x -0ls) (xtT1x-pn(s)) 0T (bs)
21k Eg 0 +
+ Esnyuy
Proof.
According Proposition C.2, one has & (no (Ts)) &g (ne(s), Hxn(s)) + + ((()))))
Hence,
(1os), = Vol (n) (x))
LK Ve [F((s) (s]l
2ik (IN 8 ((s)) (Fo(s) - F,)s))
1
= 2 (KK & (s))
(Ix 8 ((s)) 0 0+C K 1k 1xt11=-Pr(s)) pp(s))
27

---

- II @ (d(s))(s))) 0 [[0
K+1 1x - Py(s)
1
8 )b(s)
We also have the following matrix representation:
Vely (mo(s)) = ((x((s) (Fo(s) (s)s))
T = =(xx((s) (( K-1 1k - Pr(s)
CT
Proposition D.2. For any 17 E (gesian) s and ne
Gasien it holds that
2, (7. mo) 2,g, (n. 7 p,
[KK., II 2 +
+ 2,1s
II o.K mo).
The proof is straightforward and almost the same as that of Proposition C.2 if utilize the affine structure.
D.3 Proof of Proposition 3.2
Proof. Recall the definition of the distributional Bellman operator Eqn. (6) and categorical projection operator Eqn. (8), we have
(s 0) Pk (IT no (s))
X Ikk
Exx (
KK
ro + 7G - IK
E
80
8
IK
E
Pj (s1: 0)
no yo.
Ik
80
(20)
4K
j=0
=E
Pj (81; (2110))aa 0)g =
K-1
=E
9k, +
Σ 0) (gj,((ro) gk,)(ro)) 8o
1=
Hence, let W -C-T and w vec(W) (C-1 Ia) (see Appendix D.6 for their meaning). then
(s)) = (Nk(8; 0) K-1 K-1
(ro)
gxi((o) 9k.)(ro)
K-1
-E
Σ (s):0)
80 = s
9K _9KKK-1 (ro)
j=0
gj
-9K,K (1r)
K-1
E
gk(ro) +
Σ
Py (81; 0) (g)(ro) 9x(ro)) 80 =
1=0
(x())) + (G(rr) 1 gx(r))) po 1000 = 80=0
E
=F
gk((r)) + (G(ro) 1 2 (r)))
(Ix 3 (b81))
K +
K
80 =
=
(G(ro) 1 gx())) (IN @ ((>1))
80 =
28

---

+ E
gK (ro) +
(G(rr) -1k @ gk gx(r)))
1x 30
K +
K
=E (G(ro) 1 @ gx(())) @
80 = w
+
9 gg 80 =
K
+
==
or equivalently.
K Σ j=0
Pa(s) = ( (T = +
93 (m))
so =
K +
E C
9 gr(r))
8o
+
D.4 Linear-Categorical Projected Bellman Equation (Proof of Theorem 3.1)
Combining Proposition 3.1 with Proposition 3.2, know that @' is the unique solution to the following system of linear equations for θ E RR
θ - Σー 1 Egnuer
((s) (po(s)
1K
CT
K+1
T
= -ーー1 Esnny
((6))
1x
E
B(G(m)) O (o(ss))
80 =
1
Σ E
gg(ro)
80 =
CT
K+1
K
=E Egnnty
(
1
E Benuer
Ø(s)
Σ
g)(r))
11
CT
K +
j=0
or equivalently,
E -Bnwe ((s))(s)) ((((no))((
|]
K Esnph
E
((s)
Eg6(
1x
CT
which is the desired conclusion. The uniqueness and existence of the solution is guaranteed by the fact that the LHS is invertible linear transformation of O which justified by Eqn. (28).
D.5 Proof of Proposition 3.3
Proof. By the basic inequality (Lemma H.1), we only need to show
G 17 (7, II "7"""
l2 (n7 -7-))=
1 1-?
2ur (n "x + +3,m.
[gn"". II 1
1-7
Ixn" II 1 1-7
2
K(1 - 7
where we used Bellemare al. [2023, Proposition 9.18 and Eqn. (5.28)).
29

---

D.6 LinearCCateoorical Parametrization with Probability Mass Function Representation
We introduce new notations for linear-categorical parametrization with probability mass function (PMF) representation. Let W := OC = (0()), 0(1) (()),-. ................ O(K - 0(K - 2)) e
RdxK and the vectorization of W, w = vec (W) = (C-1 08 Ia) 0 E TRak We abbreviate Pmo as
Pw in this section. Then by Lemma A.2, for any E S, it holds that
Pw
W (s)) + (K 1
(21)
PMF and CDF representations are equivalent because C' is invertible.
For any The Twy
gosign 0.K'
by Proposition C.I.
2,1 (Nwy Tww,)
= Bowwr [62 (Tw, (s), Trez (5))
-i || (Pw, (5) Pw, (s))l
(22)
=Ik tr [ (W: W2) W2) Σ
||w w wallictcjex.
=Ik
hence
with the Euclidean Vix lll c)es.) if we consider the PMF representation.
the affine
space
@osign Φ.K: (2,
isometric
space
RRdd
Following the proof of Lemma D.I in Appendix D.2, can also derive the gradient when use the PMF parametrization:
V w (s), mss)) V V (Nr (s). Ikn(s))
i v |CC (Pw)() 2
21 (Ix @ ((8)) C Cp - ps))))
2 (Ix & ((s)) C (Qx)) w + K+1 lK -- pr(s))
20x [(CTC) ([()CC)(((()(( w + ( K+1 lx p @ 6b(5 }.
(0())((8)
Vw (1). = 0(x)(s) (sw(s) - ps) CTC
)n(s))
(23)
= 2(x((s) (W K+1 1x pn(3))
D.7 Stochastic Semi-Gradient Descent with Linear Function Approximation
We denote by T. the corresponding empirical distributional Bellman operator at the t-th iteration, which satisfies
(T (st) (br..) An go5
(24)
D.7.1 Probability Mass Function Representation
Consider the stochastic gradeeradient descent (SSGD) with the probability mass function (PMF) representation
W.
la (P...((s)) [T:m=--) (st)).
where V Vw stands for taking gradient w.r.t. W,- E R in the first term The, (8c) (the sec- ond term is regarding as constant, that's why we call it semi-gradient). We check that Vw 2g (The. (8t), [T: 7 (st) is an unbiased estimate of V W 2,00 (Nw.-1y T nw.-)).
Now, let compute the gradient term. By Eqn. (23), we have
W (Tww (8)), [T nw. .I (xt) ((() (Pwi (8t) - PT: me PTi (6i) T
CTC
30

---

((x((s)) (((:)) ( (
[K+1" 1x PT. prrm (s1)
1
C²C.
where PT; nw,-1 (8t) PIxx To (sr) (px ([TT (st))
k=o E IK Now, turn to
K-1
compute PT; mhr-1 (8r). According to Eqn. (8),
(
X - IKK )
pa (IT (s:)) = x~ n (1)))
1K
+
(1-
Pt + 7 - 20
R (ects)
LK
Σ Pj 11(9411;011199...8
R-1
(k,,(r.) +
Σ Pps(s41; (0j,((1)) -K (rr))
3=0
which has the same form Eqn. (20). Following the proof of Proposition 3.2 in Appendix D.3, one can show that
K
nTT (8t) = G(r.))/T +
Σ
(re).
(25)
K +1 1=0
Hence, the update scheme is
( Pp (ot) CC
W. W 2(xx
K
((s)) d W ((8 W K +1
Σ gn(
1K
CTC.
1=0
(26)
Note that our Linear-CTD (Eqn, (13)) is equivalent
K
W We-Wa-1-afis @(s() T Wi-1-o(s--)) G
(m)) -
1
Σ
-r)
-1
K+1
1=0
(27)
in the PMF representation. Compared to Eqn. (27), the SSGD (Eqn. (26)) has an additional and the step size at
C
C.
D.7.2 Cumulative Distribution Function Representation Consider the SSGD with the CDF representation
0. - a Volh (no0_1 (st), [T no (sx))
where Ve stands for taking gradient w.r.t. 0,-1 = 0T Rax K the first term 70.-, (st) (the
second term is regarding as constant). One can check that V (no. (8t), [T no. (st) an unbiased estimate of Ve 2, (Mo.... T no...).
Now, let compute the gradient term. By Eqn. (19) and Eqn. (25), we have Ve 2 (1o.-,
[ T=no--1 (8:))
((x((s))
(ot FT:no,_, (8t)
((si) -0 (11
1
T TT 1K - Pre (8)
= -(xx((s))
K+1
31

---

((x)((s)) ((or) o(s))' 0..1 ((5++1) 00-1C- ((s44))00,-11-1 (n.)
1
Egr(rr) 11x CT
K +
Hence, the update scheme is
+ 1 - ((k Fo. - FTino (8t) T
((8)) ((s)) =0,-1 - -2xx(u((s)))
- ( (00t1)TT..000 C
1 K +1
Eg((re) 1x
CT
1=0
which has the same form Linear-CTD (Eqn. (13)) with the step size 2nd K-
D.8 Linear-CTD is mean-preserving
will show that our Linear-CTD mean-prestt penn-rreerrvng, which was first discovered by Lyle al. [2019, Proposition 8). In this section, we assume the first coordinate of the feature is constant, iee, (s) = 1/vd for any E S. As stated before, we will always assume this to ensure 3 O. K can be uniquely defined.
Proposition D.3. Let Ve := (V) be the value function corresponding 0, i.c. Vo(s) (e) [G), then for any initialization of the Linear-TD parameter th. there exists (not unique) corresponding Linear-CTV parameter Oy such that Vo. Vepo furthermore, for any 1 and even number T 2 2, it holds that
(G)
Vo. Ve
Vor
VVr
Proof of Proposition D.3. Recall that the updating scheme of Linear-TD is given by
t/ - - (x)) (Ve, (st) TI - VVp.. 1 (8++1))
And the updating scheme of Linear-CTD is given by
- FT: (s1)) Fo. θ - - -ap(s.)
Let Vo == (V)) be the value function corresponding O, we have
K-1
Vo(s) =
Σ (1 Fx(s;)) 0))
k=
=IK (1 - Fo(s)) 1x
1
2(1 - 7) (K
Hence, if we take vo.1 2(1-7)'
voo for any E (,,..., d), and 0g OaK, holds that
V = Vo, 2(1 7)
1s.
We can also show that for any uh E Rd, there exists Og E [gdx K such that V = Vo. That is we need find Og such that for any S.
IK Σ (b))T0((k) 2(1 7) k=
1
T
We can take O, satisfying the following equations make the above equation hold
K-1
Vd
IK
lo 1)
2(1 (1) 2u,1,,
32

---

and
LK
Oo(k)-1
-10
k=0
Furthermore, for any & 2 1, we have for any E E.S
1
Vo. (s)
L (x
2(1 - 7)
1
[K(((s) (0.-1 a((())) (Fo. ( ') 1x
2(1 -7)
= -Vo._. (8) ouk (((s)) (Fo. FTino, ( 1x
Ve Tbb
= ( - a0(())) (V (8:) - - VV (8++1))
=V (8) a (o((s)) (Vei-((8s) re 7 VV (8(+1))
We need check that, if Vo.-1 V holds that
LK (FT (s+) Fo.., (st) T 1x Von. Tx (s+
which is the direct corollary of the following fact
LHS Val. (8t) Ex~II (br N0. ( Ves, (8t) - - -VVo (81+1);
by Lemma D.2. And we can obtain Vo,- = Vey by using the facts that gosign Φ. is a linear space and taking expectation is a linear operator.
is an affine space, %
Lemma D.2. For any gosgn, holds that
Ex Ex~mgu[X]
Proof. By Eqn. (8), Hg Peiin
uniquely identified with a vector pa (pn( K-1 E BK
where
Pk(v)= = 10,11-7)-1 - (1 - 1( - ((dd))
Hence, for any x E [0, - (1)1-))111 -1, we define Fib max(y E (
S y). then
(2 2x)/ 1)+ =Ith (1 x --==m + (I'h + (x) (1
(1- TTb LK
Σ
1K
K
k=
=Iib + LK (1
Eib + LK
iK
=I,
therefore,
= Σ xb I (0.(1-9)-1)
Ex Ex~M
(1- 11-1(x-11) (x--xx//10 ((xx)
k=0
k
Σ
xu(1 - / (dx))
10.(1-7) 11 k=
-/p.11-)) (vx(
[Ex=u[X].
33

---

E Omitted Results and Proofs in Section 4
E.1 Proof of Lemma 5.1
Proof. By Lemma H.1 and Eqn. (22), we have
(C(0))² = W?,, +,,
(no. 17-)
1
- 1 - y L2
170 (700770-)
KK
tr (6-0") 0-) - (
(6- (0-01)
y
K(1 - 7)2 10-00 1 0'
We only need to show that
Ix & EΦ & (1- vマp A (Ix (Ix A ((x
4 ATA
< (1 7)2Amin
(28)
or equivalently,
(Ix 1x020'0 A Ix 2 x A [k 8 x = (1 (1-v7)² Lak.
Recall
A= (Ix & ΣΦ) - Ear, [(GG())-)) 8 (+()))(6))))
then for any 0 E R with ||0|| 1,
IK & Σ # A
[I & E3 A Ix @ 0
θ
IK & E +) A [K o x oo|²
(CGG))--) 30 (x+((0))())) (-)0)))))))))))))))))))))))))) イ 1²
0 - Eana
2 |Berro [(C@()C-)) @ (Z (X)000)))) ) 0 01)²
It suffices to show that
|eerre ((G())-1)) 3 (I (1))(())) Σ3 #1 )
SV
(29)
For brevity, we abbreviate C&(r) as ( = (wy (r) i= E R Thus, it suffices to show that
| Y @ (x $ (()))(s))) T 1
For any vectors w- (w(O)T, ..... ww( - )) and u (v(0) ..... v(K-111 1) in Rak we define the corresponding matrices W (w(0), w(K 1)) and V (v0(0), (.. v(( - 1)) in d then ||l||| = || F /tr (W W) and ||||| VVlF vtr //t V). With these notations, we have
(r)) & E (s)))(()))))))))))))))))))))))))
w E Ea ((r) s E ()))())))))) (
sup |||||=||||=
K
(1(())((i) T ( t
sup
E Esr
x(1)
||||| =|
:===1
it is easy to check that
T Σ wii)
1
= ( (
34

---

hence
E Ba Y(r) 2 Dx (()))(())))) Ex
Σ (sy(r) W ((s))((() ( + 13
sup
Barry
|
Ir (((())) ΣΦ (ss')((()))))) w)]
sup
Barral
|W]|==I|VI|
sup
B Barre
((5) E WY xbs5))
|
W E ((()) ||())|| v Σ $ <(5>)||
sup
Eg Rr
|
sup
B
W T Σ (66))
[W
Vil =1
Σ : Φ ||] E.y ||T /|
sup
Eg
W
|W V
- v7 || -VII
sup
E.
(b)) Σ
WW Σ [(}) E. ( VV Φ (()0))
(WW ES E, (0(0)(()) ir E (+()))(()))))]
=
sup
[[ v
sup
(WW tr (VVT)
|W||,
VII VI,
sup
|
[|]==
= V7.
where we used (r))ll S v for any r E [o, 1] by Lemma H.3, and inequality. To summarize, we have shown the desired result
(IK 1K E A [K @ E A A(Ix @ x > (1 (1-v7)² Iak.
E.2 Proof of Lemma 5.2
Proof. For simplicity, we omit t the random variables, for example, we use A to denote a random matrix with the same distribution as A,. In addition, omit the subseripts in the expectation, the involving random variables are s ~ Ma, a ~ π (),))))))) (r, ~ P(- s,a).
Bounding CA.
By Lemma A.3,
|||| < [ ( + (| (CG(M)-))) @ (d(s))(s)T)
=|b((>)] [+)))) |
<1+ V7.
where we used Lemma H.3. Hence, CA S 2(1 + VT).
35

---

Bounding C..
By Eqn. (32),
| = (0*)² A A0*
22 |1011 (((())))))) +
22 + 7) sup
10 12 ((((()))())
#ES
22 + 7) sup
[l((s)) ||*1
sES
22 + +7) |
Hence
||40|| SV2(1 |1111
As for |[b]]
K
|b||
C
Σ (r(r) - 1x
8 (((5)
K +1
1=0
k
C
Σ gj(r) - 1x
[
(30)
K +1
1=
K
C
99 (r) - 1K 1=0
Σ
K +1
By Proposition C.3 with n E Pkan satisfying (7()) = v for all E S, where v
(K + 1)_1 Ex === 8z. is the discrete uniform distribution, can derive that, for any r [0, 1]
and s' E S, it holds that
Egi(r) gr( - 1x
K+1 1x
1
-G(r) (nn('')
K + 1x
K+1
1=U
1
P(bra))
1k.
K +
Hence,
|b||
C Σ
(r) 1K
K +1
D(brn)
K 1x
VLK le (bIn) (v),
< K(1 - 7)6 (1bm)) (v),
<3v 77),
where used the orthogonal la ((br)) (v), (Lemma H.4).
decomposition (Proposition C.2) and upper bound for
In summary,
|e|| = A|A bll
< <||A0' || + |b|||
S 2 + + 3v (1 7))
Hence, C. S 2((++7)
10 +3 K (1 (1-7).
36

---

Bounding tr (Ee).
ur (ree)) =B [ - -b]o]
=2(0") T "B[AA [AATAJO
AJ 0" +2 [bb
By Eqn. (33),
(0*)* E [[ A 0 22 +7+ 110
And by Lemma H.4,
I b bj <9K(1 - -7)
To summarize,
tr (Z.) <4(1+)) || |1 IK & + +18K(1 - 7)
E.3 Proof of Lemma 5.3
Proof. For simplicity, we the same abbreviations in Appendix E.2. As in the proof of [Lemma
Samsonov et al., 2024b]. we only need to show that, for any E N. a E (0. -1 (3)p))))). it holds that
{ [Lak- [aA) (Lak - (Iak-aA)
-aA)]"} & Lak 1 (ak-2o0((-))K000
(5(1-2 2ap(1 - Am)n
(1-1
Lak
Let B == A + A aA T A which satisfies (Lak aA) (Iak aA) Lak aBB. To give an upper bound of E [Lak - =, it suffices to show that
17 Es- Vp (2.3....).
E B > (1 I k & Es.
B
4 K
16
if we take E (0, (1 (2 +)))
Given these results, we have, when E (0. (1
(38p)) it holds that
E [ aBB) xI apE B +
Σ
E
B
1=2
$1-
ap(1 - /7)
17
(4ap)'
IK @ Es
16
1=2
ap(1 - v7) 17a2p²
1K ② Es. 4op
1
31
ap / & Es.
2
Lower Bound of [B). To show B - Vz Eg, first show that E A AT AT] x
2(1-v7) & Eg. which is equivalent (Ix - E [ A 1] (Ix @ Ex) x 2(1-v7).
where
E A - A
2 (Ix & Ex) - E (CG())=)) -
(cā(r)))- @ (
Then, for any BicK with ||0|| 1,
0 (IK & EΦ)
E
A + A
(IK 2 Σ6)
2 T0
(Ix
E [Cā(r))))) (6())((5))) (Ix -
2 2 ||Eenrd (CG @ (1)0(1))()))))
21 v7),
where we used the result Eqn. (29).
37

---

Next, we give an upper bound for E [ATA] AJ. we need to compute the following terms: by Lemma A.4,
[Ix @ (o(s))(s)
2122 =Ix (d(s)o(s)
=1b(())| IN @ (0()))(s))) xIx @ (o(s))(s))
(31)
Hence
E [Tx @ (+()))()) l @ EΦ.
And by Lemma A.4 and Lemma H.3,
[(Gā()) (0(9))))))) ((C@N)))) o (0())(())))
=n ( ((5))(5)
= =I10(+)) ( (6())(())
< ||c@@nc-1|1 In (R (
< @ (0()))(())))
To summarize, by the basic inequality (Br B₂) (B, - B2) x (B] B1 + B3 B2), we have
A xII (b()))(8)) + +7
(32)
and, after taking expectation,
E A A] <2 + IR E..
(33)
Combining these together, we obtain
E B [ /7) - a(1
Ix 2 Es (1- V @ Es.
we take a E (0, (7)/(((1+
Upper Bound of E B Because 2 is always PSD, we have the following upper bound
B² B 2 B
We first give almost-sure upper bound for B By Lemma 5.2, ||||| S 1+ Vyi And by Eqn. (32),
AA |AA
A|l
<2 Ix @ (o(s)Φ(s) +
<2 ())))((s)
+00(s))0())
(34)
<2(1 +))
Hence,
BB| = || + A T aA All
<2 |A| + A ATAI| All
(35)
<2(1 Vx) 2a(1 +7)
<4.
because (0, (2 + +)))
Now, aim to give upper bound for E [
B² (A+ A T aA T A)²
<11 + 3) (A AT)² 2 + (1+ (1+ a² A)²
22(1 + 3) (A A+ AAT) + (A A)?
38

---

where we used the fact that (B, + B²) (1+ B) + (1 (1+ B for any symmetric matrices 2 -B31-33TBB) & 0, 3 E (0. 1) to be
B1, B2, since 3B$ + 3 1 B3 B, B,B1
determined; and the fact that A² (4T)² A AA since the square of the skew-symmetric
matrix is negative semi-definite
( * 0. By Eqn. (34) and Eqn. (33), we have
2
|| A S 2(1 7);
[A] A] <2 (1 + 7) Ix @ ΣΦ,
(36)
thus, by E (1, - (2) + +))), it holds that
a² (A T A) 21 + 7) IK @ Eg < (1- (1- @ Eg.
(37)
As for EE |AAT, by the basic inequality (Bi - B2) (Bi B₂)
ক BB + BBBB)
have
AA T
[IK @ (o(s)((s) [)] - [(GGm))))
@ ) (b(())(())}}
{LA (0()))(s) []] [[CG (0(6))()))
(
x2 [I @ (+())(+)) ))² + 2 ( @ (+())(())) [(G(())-) @ (((0))())))
By Eqn. (31), we have
[I & (p(s))(s)))
<Ix +Ix (0(s)
And by Lemma H.3,
(0(3))())))
(666)66)))
= =(GGN))))= ((6666())))))))))
= I10((11 (CāN)C-1-TTT (6())())
||cG@n--l| Ik (8(s))(s))
< < @ (o())((s)))
To summarize, we have and after taking expectation
AAT <2 + Ik & (p(s))(s)))
E [[AA] ] <2 (1 + 7) In & Es.
(38)
By putting everything together (Eqn. (36), Eqn. (37), Eqn. (38)), we have
B 22( + +.3) (A A + AA + (1+
A)² < [8(1 + (1 3) + (1 ( (1-)))))))))) Ix 8 Es
(17 + (17+97-10/7 Ix & Σs
171k @ E.,
where we take 3 88(+7))
Therefore, by Eqn. (35).
Bº < B 1- P-2 B²
<4p 24p-217K @ E.
17
IIx s E..
16
39

---

E.4 LP Convergence
Theorem E.1 (L" Convergence). For any K 2 (1 (1--))_1pp and a E (0, / log T),, it holds that
100 1x +1
E [[C (0+))"
(K(1~2)
v exp + op
VT
(1- (1-7)VAmil
7)
10 (0"1IX88000 +1
VK(l-7)
1 Vap
1
7
(1 - 7) Aumin
1 (1 (1- 7 )
1/2
T
va(1 - Am
Va
v - 71--))minn
p
VK -
100 0-
Proof. Combining Lemma 5.1, Lemma 5.2 and Lemma 5.3 with [Theorem 2 Samsonov et al., 2024b].
have E
[[C (Or))
ptr (Ee)
CaVap
CaC.op /tr (Ec)
(1+ CA)CP
14
/ - (()1-))Ammi))
T
Va
T
+
(r
1
+ (1 aa)
1
+
CAP VaaT)
|10 0*1l
VaT
Vap
aT
(100' || + VK (1 7)
10 10 1I + VK(11 (K) -7)
Vap + ap
[1
VK(1 77)
VT
((1 7)
T
1
IIN + VK(1-1)
1
+
1
V(1 (1))mim
Vap
+11
1
a(( - (5))min) 7//2
+
110 0-1
2
aT
v [x(1 - -7)>mmmn T
VP VK(1-)) |0 IN +1
v axp + ap
1
VT
(1 (1-18 Amiin
7)
110 IINDE +1
(k()-))
1
+
1
+
T
(1- Amin
Vap
+ (1
fa(l1 (7) T/2
1
p
+
VK(1- 7) 110u -
va(1 - Amia
Va
v(I
-
7)
where used the fact that ||0*|| S (Amin)
1/2 |00
1
E.5 Convergence Results for SSGD with the PMF Representation
In this section, we present the counterparts of Lemma 5.1, Lemma 5.2, Lemma 5.3 and Theorem 4. for stochastic eemiigrrdient descent (SSGD) with the probability mass function (PMF) representation. These results will additionally depend on K. The additional K Kddependent terms arise because the condition number of C C scales with X² (Lemma H.2). These terms are inevitable within our theoretical framework. The proofs of these results require only minor modifications to the original proofs, and we omit them for brevity.
In fact, in Appendix G, we validate some theoretical results through numerical experiments. To be concrete, we find that empirically, as K increases, ensure convergence, the step size of the vanilla algorithm in [Bellemare et al., 2023, Section 9.6] indeed needs to decay at rate of K-2, In contrast, the step size of our Linear-CTD does not need to be adjusted when K increases. Moreover, we find that Linear-CTD empirically consistently outperforms the vamilla algorithm under different K.
40

---

Recall Eqn. (26), the updating scheme of the algorithm is
Wi < W acp(s))
Pw:-, (st) pw: (5:) C
- cpp(s))
((x:) T W ((8111)
(r-1G"(rr)) " 1
-g 1x
c²c,
K +
1=0
which is equivalent
W.C WW--ic"-apla. ((si) ((si) Wi-1C -(-1cc= W...C
Eos cT cc
K+1
here we drop the additional 2 in the step size for brevity. Letting PPMF, = W.C be the CDF parameter, the algorithm becomes
K
PPM 1PMP;----0((6 ((8)) OPMF,C1 ((8)-1)
( I K+1
E0 a ca
(39)
Here, we add the subscript PMF to the original notations to indicate the difference. Then, the algorithm corresponds to the following linear system for 0 E RdK
O 0 bYMP;
where
APMF
[(CCT) & ΣΦ] - -B...' (CC" & (0(5))()))
{ C ( K 1x
bPMF = K Esr
-
8 ((s)
Compared to our Sinear-CTD (Eqn. (16)), this algorithm has an additional matrix CCT with the condition number of order K2 (see Lemma H.2).
Now, we are ready to state the theoretical results for the algorithm.
Lemma E.1. For any Rak it holds that
(()) S
ApMP (0 0-)|-
(40)
VK(1 - >/2 /mmin
Amin
Lemma E. achieves the same order of bound as prior results for Linear-CTI CTD (Lemma 5.1), as the minimum eigenvalue of CCT remains Q(11 (Lemma H.2). However, from the numerical experiments (Figure 6) in Appendix G.2 we observe that after substituting O, into 0, as grows, the RHS grows with K, while the LHS remains almost unchanged. This might be because when the matrix CCT acts on the relevant random vectors, the stretching coefficient (i.e. |CCTr))) ||æ||] for some vector x)) is usually order K2 rather than a constant order. For example, consider the case where the matrix CC acts random vector X that follows uniform distribution over the surface of unit sphere [|] = 1). Since the k-th largest eigenvalue of the matrix CCT is of order k.2, by Hanson- Wright inequality [Vershynin, 2018, Theorem 6.2.1], we have CC X|| of order K² with high probability.
Lemma E.2. It holds that
CA,PMP S K2CA, Ce,PMP S K2Co, Ir (r(2.,MF) S K ir (E.).
Lemma E.2 introduces additional factor of K² (or K*) compared to previous results for Linear-CTD (Lemma 5.2). since the maximum eigenvalue of CC is order K²,
Lemma E.3. For any p 2, let apmr (1- >mii and a ~ (1 vT)/(p33) PMF p S 1 1/2). Then for any a (0.a PMF). u E RdK and t E
(a)
L ||]
S S "1aapp))) ||
41

---

Concepts
Linear-TD
Linear-GTD
Parametrization
(+) (18) V
Fite: ((a) ((A)
Bollman Operator
(2 V((e) Elru V(a.)
( R)te)
Projection Operator
1/3 V.
2 E Panne ((())(())
1/6
1 9)
(ols)( Fols)
Projected Bellman Equation
11. VV.TV...... Vi. [See Egn.
11
7/-. See Egn. (12)
Update Rule
the e- d 1 - -n0())) (0(s)) (019111))
oh 1 ((a,)((a.) -(0-))
[See Eqg. (13)
A Key Quantity
(d(s))((s))) I [CG ((r.)
(0(o,)
CG (rr (r,)C with spectral /
b,
re((),)
R [C(2] 2,fre) 1x d(a))
Key Quamtity in
K 4/2 7) |C ( 1xl1
Measure of Emror Approximation Firor
Basg. (V(s)
Wi (7) 7j (Cange W Fins). (SD)
Propoction 3.3
Sample Complexity
7P
Table 2: Comparison between Linear- and Linear-CTD.
As before, in this lemma, 4PMF does not depend on K because the minimum eigenvalue of CC is (11), and a PMF scales with K because the maximum eigenvalue CC is of order K?,
Theorem E.2. For any K 2 (1 11)) and ap.0 P.
), holds that
K2 (110* llv, 11
ak²
K3 (1101l1, +1)
EV''(C((PMETT))
1 1
VT(1 - Amin
(1 7)mmin
TVak? - 7) Amain
{(ak3)
(i)Amin)
K2
K
| |00mMr - 0" 0'11v3
Tvak? 7)) Amin
Vak?
((I7))
This theorem for the PMF version algorithm yields an upper bound that is K times looser than Theorem 4.1 for our Linear-CTD. The appearance of the K3 factor is due the fact that the condition number of the redundant matrix CCT is of order K², This factor is unavoidable within our theoretical analysis framework.
However, from the numerical experiments (Table and Table 4) in Appendix G.2, we can only observe that our Linear-CTD consistently outperforms the PMF version algorithm under different values of K, but the performance gap does not increase significantly when K becomes larger as predicted by Theorem 4.1 and Theorem E.2. The reason for this might be, discussed after Lemma E.1: 1: in the experimental environment we have set, when the matrix CC acts on the vectors it encountered, the stretching coefficient is usually of order K² rather than constant order. See the numerical experiments (Figure 6) in Appendix G.2 for some evidence.
F
Comparison between Linear-TD and Linear-CTD
To further improve the readability of the paper, we compare various concepts and results between Linear-TD and Linear-CTi in Table 2.
G Numerical Experiment
In this appendix, we validate the proposed Linear-CTi algorithm (Eqn. (13)) with numerical experiments, and show its advantage over the baseline algorithm, stochastic semi-gradient descent (SSGD) with the probability mass function (PMF) representation (Eqn. (39)).
To empirically evaluate our Linear-CTD algorithm, we consider a 3-state MDP with 7 = 0.75. When the number of states is finite, denote by Φ = ((b()) E R the feature matrix. Here, we set the feature matrix to be full-rank matrix in R x3 The following experiments share zero initialization On 0 with max 500000 and batch size=25.
All of the experiments are conducted on a server with 4 NVIDIA RTX 4090 GPUs and Intel(R) Xeon(R) Gold 6132 CPU @ 2.60GHz.
42

---

G.1 Empirical Convergence of Linear-CTD
We employ the Linear-CTD algorithm in the above environment have the following convergence
results in Figure 2. This figure shows the negative logarithm of K 0
0*
11 = (1
10
logl e
9
8
7
K=1
6
K=2
K=5
5
K = 50
4
K 500
0
10000
20000 Iterations
30000
40000
Figure 2. Convergence results under varying K for our Linear-CTD algorithm with step size = 0.01. These curves exhibit similar trends, demonstrating algorithm's robustness across different K values.
1) (no no-) along iterations. We observe that our Linear-CTD algorithm can converge for different values of K when we set the step size 0.01.
G.2 Comparison with SSGD with the PMF Representation
First, we repeat the same experiment as in the previous section for the baseline algorithm, SSGD with the PMF representation. The experimental results in Figure demonstrate that when the baseline algorithm uses fixed step size
0.01, it does not converge when K is large (K > 44). The results in Figure and Figure verify the advantage of our Linear-CTD over the baseline algorithm 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-CID does not need be adjusted when K increases (Lemma 5.3).
PMF,(K) that ensures the convergence of the baseline Next, will verify that the maximum step size a algorithm scales with K as predicted in Lemma E.3. Then we will compare the convergence rate of the baseline algorithm with that of our Linear-CTi algorithm.
In Figure 4, we employ the baseline algorithm with different step sizes under fixed K = 150, and we find that the baseline algorithm converges when the step size does not exceed 8.6e-4, and it does not PMF,(150)
converge when the step size exceeds 8.7c-4. This indicates that doo this environment, providing good approximation of axo
[8.6c-4, 8.7...77444 in
PMF,(150)
We repeat the above experiments under varying K, looking for step size that PMF, (K)) and a step size that leads to divergence (an upper bound PMF, (K))
convergence
(a lower bound of acc such that the two step sizes as close as possible and thereby we can get good approximation of a'00
arso
(K) The results are summarized in Table 3.
PMF,(K) provided in Table 3 to perform a quadratic In Figure 5, we use the approximate values of aocc
function fitting of 1/a PMF,(K) with respect to K. We find that a PMF,(K) indeed approximately scales with K which verifies theoretical result (Lemma E.3).
To compare the statistical efficiency of our Linear-CTD algorithm and the baseline algorithm, in Table 3, we also report the number of iterations required for the error reach below 2e-6 when the
43

---

10
5
0
K=1 1
K= 2
5
-5
K 44
K 45
-10
10000
20000
30000
40000
Iterations
Figure 3. Convergence results under varying K for the baseline algorithm, SSGD with the PMF representation with step size a = 0.01. We remark that when K 45, the program reports errors of inf and nan. In contrast to results of Linear- in Figure 2, the baseline algorithm no longer converges when K is large (K" 44).
16
14
logg lla. lke
12
10
a = 6.0e
8.6e 4
a 8.7e 4
a 8.8e
10000
20000
30000
40000
Iterations
Figure 4. Convergence results with different step sizes for the baseline algorithm, SSGD with the PMF representation under fixed K 150. We remark that when take 8.8e-4, the program reports errors of inf and nan. The baseline algorithm converges when the step size does not exceed 4. and it does not converge when the step size exceeds 8.7e-4. Therefore, PMF,(150) E [8.6e 8.7e-4) in this environment.
step size satisfies a 0.2ccc PMPPPM PMF,(K) In addition, we present the parallel results of Linear-CTD in Table 4. In Table 4, we find that the value of areo for our Linear-CTD algorithm is much larger than d'oc comparing the Iterations columns in Table and Table 4, we find that the sample complexity of our Linear-CTI does not increase significantly with the growth of K, and Linear-CTD empirically consistently outperforms the baseline algorithm under different K.
PMF,(K) and it does not decrease significantly with the growth of K. Moreover, by
However, the performance gap does not increase significantly as expected when K increases as predicted by Theorem 4.1 and Theorem E.2. The reason for this might be that, as discussed after Lemma E.I. I, in the experimental environment have set, when the matrix CC acts on the vectors
44

---

K
||x||| for some vector æ) is usually of order Lower Bound of PMF.(K)
Upper Bound of PMF,(K)
Iterations
30
2.1e-2
2.2e-2
37245
45
9e-3
9.5e-3
39262
75
3.4e-3
3.5e-3
38286
105
1.75e-3
1.8e-3
38123
150
8.6e-4
8.7e-4
38556
225
3.8c-4
3.9e-4
38317
300
2.1e-4
2.2e-4
38999
375
1.35e-4
1.4e-4
38674
450
9,5c-5
9.8e-5
38506
Table 3. Lower and upper bounds of the maximum step size ax PMF,(K) that ensures the convergence under varying K for the baseline algorithm, SSGD with the PMF representation. The bounds determined using the method that in Figure 4. The Iterations column refers the number of iterations required for the reach below 2e when the step size satisfies 0.2
PMF,(X)
Approximate values of 1/a
10000
Quadratic regression curve
8000
1/a PMF,(K
6000
4000
2000
y 0.0507x + 0.3034x 17.4170
0
100
200
300
400
K
Figure 5. The approximate values of of maximum step sizes 1/ PMF.(K) under varying K. Here we PMM.(K) provided in Table as an approximation take the average of the upper and lower bounds of Ckoc PMF.(K) and perform quadratic regression of 1/ PMP.(K) on K. This achieves mean squared a
error 425.85 and R² of 0.99996, which indicates that 1/ PMP.(K) indeed grows quadratically with respect to K, aligning with our theoretical results (Lemma E.3).
it encountered, the stretching coefficient (i.e., |CC x 2 rather than constant order.
We verify this conjecture through the following experiment. We focus on the LHS and RHS of Eqn. (40) in Lemma E.I:
(()) S VR(1 2 Ammin
ApMP (0-0*)1)
where
APMF = [(CCT) & Eg] E (CCT)))) @ (0()))))))))))
In our theoretical analysis, we first give an upper bound of the RHS, and then apply Lemma E. bound the loss function C(0) in the LHS with the RHS. However, since the minimum eigenvalue
45

---

K
Lower Bound of dooo
(K)
Upper Bound of (K)
Iterations
30
1.65
1.7
17908
45
1.65
1.7
17925
75
1.65
1.7
17942
105
1.6
1,65
18623
150
1.5
1.55
21947
225
1,55
1.6
20890
300
1.55
1.6
20890
375
1,55
1.65
20595
450
1.5
1.55
21947
Table 4. Lower and upper bounds of the maximum step size di that ensures the convergence under varying K for our Linear-CTi CTD. The bounds determined using the method that in Figure 4. The Iterations column refers the number of iterations required for the error reach below 2e-6 the step size satisfies 0.2c (K)
(K)
of the matrix CC in the RHS is only of constant order, unable to have term of 1 K? in the RHS. Therefore, our conjecture can be verified by checking whether the bound provided in Lemma 1. is tight in this environment, which is presented in Figure 6. The left sub-graph of Figure 6 corresponds to the LHS, and the right sub-graph corresponds to the RHS. We omit the constants that independent of K. From Figure 6, we can find that the LHS remains almost unchanged under different K, but the RHS increases as K becomes larger. This indicates that the stretching coefficient of the matrix CC that frequently encounters during the iterative process grows with K rather than remaining a constant order. A similar analysis also holds for 6PMF in Lemma E.3, and we omit it for brevity. These factors result the performance gap between our Linear-CTi algorithm and the baseline algorithm not increasing significantly when K becomes larger, as predicted by Theorem 4.1 and Theorem E.2.
-log Am -'
-2
-4
-6
-8
K
K=45
K =
10
6=
75
-12
1500017500
15000 17500
Iterations
Iterations
Figure 6. LHS and RHS of Eqn. (40) Lemma E.1 under varying K. The left sub-graph corresponds the I.H,, and the right sub-graph corresponds the RHS. We omit the constants that independent K. We can that LHS remains almost unchanged under different K. the RHS increases K becomes larger, indicating that the stretching coefficient of the matrix CC that we frequently encounters during the iterative process grows with K rather than remaining constant order.
H Other Technical Lemmas
Lemma H.1. For any E Spignn we have (1))) S 1
-í (0,,,))111)))
46

---

Proof. By Cauchy-Schwarz inequality.
(1(,
Frn (x)
(x)(2)d
1/-
SV 1.
12dry
[F)((x) (2)
1
V1 U2).
Lemma H.2. Let C
E ER be the matrix defined in Eqn. (11), it holds that the eigenvalues of CTC are 1/(4cos/ (KT + +))) k E [K], and thus
|| ClI
S K2
|[TTo]_1| cos²
S 4.
4sin 1K+2
2K
Proof. One can check that
K
K -
2 17 2
K-1
c²
:
2
2
2 1
1
1
1
-1
0
-1
2
-1
0
-1
2
0
(CT
Then, one can work with the the inverse of CTC and calculate its singular values by induction, which has similar forms to the analysis of Toeplitz's matrix. See Godsil [1985] for more details. O
10.1], it holds that |cG())--1| s
Lemma H.3. For any
vT and
[(TT 1/2 (C)) - 1/2
S VT.
Proof. One can check that
07
0
0
C-1=
is clear that
[(TT)) c) 12 (1)
SV + (c C OF *
G <C
- (r) <7IK
c)()) S
By Lemma 1.2 and upper bound on the spectral norm (Riesz-Thorin interpolation theorem) [Serre, 2002, Theorem 7.3], we obtain that
|cā < || cā@n)-1| $ V1. V/.
47

---

Lemma H.4. Suppose K 2 (1-2)-1, = (K+1)-1 -1 Ek=o then for any E [0, 1]. it holds that
k= dze is the discrete uniform distribution.
82 ((0r,)) 2 =3)
Proof. Let v be the continuous uniform distribution on [0. (1 -7) + ux]. we consider the follow- ing decomposition
&2 (v, (bra) (v)) S l2 (v, + 82 (0, (Dr,)) (0)) + 82 ((br))+ # (8), (D (brx)#(())
By definition, we have
(2(1)) -V (K +1)+ 1/- PK (10-))
K 2
1
dx
K
1 V 3K(K + +1)(1-7)
KVI->
By the contraction property, we have
V
Ee (brr)) ) (v), ( (v)() S
KVT
We only need to bound Eg (i, (br,)) # (D)). We can find that (br,)) # (7) the continuous uniform
distribution on [r, + 7-K + 71 7) ")- and the upper bound is less than the upper bound of v,
namely, r + yek + (1 - 7)- $ (1 - 7)
PIK- < (1 (1)-7)-1 + LK. Hence
(2, (6 )+ (B))
/-''
(1 (1-1) 7)
K x+1% E
K
2
dr + 1
K
I --r r )}²
dx
K+1
7
(1-7)-1 -1 +Ck
11 - 7) (1) K 2
+
K + x
dx
(+-0)Ct((111))1
(1 K2,
(1-7)??2
(1 - K2 ( (K+1 -r) 31 ) (1 11-7) 7) [K1 -r)²
3(K + +12
3(K + 1)2
3(K + 1)2
3(K + 1)2
5(1~2) + (1 - 7)h
=1 7.
To summarize, we have
1
+ 7+ V7 < 23 KV
8a (v, (br.) # (v)) 5
K KVI
where we used the assumption K 2 (1 -1-7) 7.
I Analysis of the Categorical Projected Bellman Matrix
Recall that G(r) = G(r) - 1 03 gk (r). We extend the definition in Theorem 3.1 and let gj,x = he ((r + yay- - rk)/LK) = h(r/ + 7j kk) for j, h E 1,,,,..., ,K} where h(x)) = (1 - [=])
Lemma I.1. For any [0, 1] and any A E 1. KK. in gr(r) there is either only one nonzero entry or two adjacent nonzero entries.
Proof. It is clear that h(x) > 0 = -1 -1 < 1. Let (r) = min (k: : (r)(r) > 0}, then (r) = min(k: {k: trj is due to
+nj - k < <1 = min : 0 r/uk +7j - k < 1). The existence of (r)
T/IK + 1j - K < 1/uk + - <1-)
K K=0<1.
Let (r) = 7/LK + - (r)) E 10, Then (r) = h(a,(T)) 1 ar(r) and (1,,())+(T)) = h(a,(r)- 1) (r) are the only entries that can be nonzeros.
O
48

---

The following results are immediate corollaries. Corollary I.1.
0,
for < < ky(1),
Σ (r) R=
= - as(r),
for = ky(r), for ky <is K.
1,
Corollary 1.2.
kr+)(r) (r))));
fa,(r)
(r) +1, (((r))1-9)
As a result,
[a,(r) +. 11-
(r+1(r) (a,()) +7-1, fa,(()>>-9-
Lemma 1.2. All entries in CG(r)C-1 are non-negative.
| and
[C ||, 1.
Proof. By definition the entries of G(r) are (G(r)),, (r.((r) (r, for j, E {0,1, KK11}
Using the previous corollaries, through direct calculation we have that if kz+1(r) = kr)((),
0,
for < kz(r),
(C = Σ (g)(() Σ (r) = a)41(()) - a,(r),
for = (r),, for kr(r) <i i < K.
k=0
k=
0,
And if kfti(r) = kr)rr) +
0g for 0 < i<< ky(),,
1-aa(r), a)t1(r),
for i = ky(r), for i = kr)z(r), for (r) < <.
(Cē(r)C))),, () Σ (7)
k=0
k=0
As result, all entries in CG(r)C-1 is non-negative. Moreover, the sum of each column and C is since
K-1 (C (r) =
if (r)=
[1 ar(r) + (z++(()) = =7 if kg)r
Moreover, the row of CG is
K-1 (C
90.m (0 Σ (r. S 0 1.
1=
m==
m==
Thus, it holds that (G(()C-1 51
49


Reviews:
[{"review_id": "7zNPu9ayl7", "reviewer": "Reviewer_8d43", "summary": "This work derives finite-sample rates for distributional TD learning algorithm with linear function approximation. Specifically, the work formulates a projected Bellman operator related to the distributional RL setup and builds a stochastic approximation algorithm to find its fixed point. Their main result concerns the sample complexity of  this algorithm in the generative and the Markovian setting. It is shown that the algorithm needs $O(\\epsilon^{-2})$ samples to get an $\\epsilon$-accurate solution estimate. This sample complexity matches that of classic TD learning.", "questions": "1. Why have you added $(k + 1)/(K + 1)$ in (10)? Please add detail on your claim that it is for normalization.\n\n2. What are the limitations of working with signed measures? 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?\n\n3. 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?\n\nBertsekas, D.P., 2011. Approximate policy iteration: A survey and some new methods. Journal of Control Theory and Applications, 9(3), pp.310-335.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**:\n\n1. Paper is well written and easy to read. \n2. The paper is able to analyze the challenging distributional TD learning algorithm with linear functiona approximation. Distributional RL has been a promising direction in recent literature and this work should enable their \n\n**Weaknesses**: \n\n1. The proposed theory and algorithms apply to the space of signed measures, rather than usual unsigned measures.", "was_revised": true, "final_justification": "The paper has analyzed distributional TD learning with linear function approximation. I believe the work should be of interest to researchers working on disrtributional RL.\n\nHence, I recommend a clear accept.", "rebuttal": "Thank you for your valuable review and suggestions. \nWe are very glad to know that\nyou think our paper is well-written, and our results are helpful for subsequent research in distributional RL.\n\nRegarding the weaknesses and questions, we provide the following detailed responses:\n\n> 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?\n\nWe are grateful to you for raising this issue.\nWe 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.\nThe non-probability signed measure does induce some limitations.\nFor instance, this makes some common statistical functionals such as quantiles and CVaR ill-defined, thereby failing to guarantee their convergence. \nWe will provide a more detailed discussion in the revised version.\n\nHowever, 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.\nThus, our finite-sample theory built on signed measure spaces still finds applications in many problems. \nFor 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.\n\nTo 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]. \nHowever, these parameterizations leads to nonlinear algorithms, which may prevent us from proving finite-sample convergence comparable to that of Linear-TD.\nCompared 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. \nRegarding the finite-sample analysis of the aforementioned nonlinear algorithms, we leave this as future work.\nWe hope this explanation is helpful.\n\n> Question 1: Why have you added $(k+1)/(K+1)$ in (10)? Please add detail on your claim that it is for normalization.\n\nThank you for your valuable questions and suggestions. \nThis term serves as a normalizer.\nBelow we will provide a detailed explanation and include it in the revised version.\n\nAs 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$.\nFirst, we explain why normalization is necessary.\nAlthough we work with signed measures, it remains necessary to ensure the total mass is normalized to $1$.\nThe reasons are from both algorithmic and theoretical considerations.\nFrom an algorithmic perspective, without normalization, the difference between our estimator and the true return distribution would not be a signed measure with mass zero. \nThis would make our estimator less interpretable and lead to issues in defining many important quantities. \nFor 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$. \nFurthermore, 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.\nIn this sense, the algorithm will not be a natural extension of Linear-TD and tabular CTD.\nFrom a theoretical perspective, tight sample complexity bounds can only be achieved after introducing normalization. \nWithout 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.\n\nNext, we explain the impact of choosing different normalizers (distributions used for normalization) and why we choose the discrete uniform distribution $\\nu$ as the normalizer. \nAs 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. \nFortunately, such undesirable properties would not appear in the tabular case.\nIn 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. \nIn 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). \nBased on the above discussion, any distribution could be used for normalization in practice.\n\nAnd we choose the uniform distribution as the normalizer only for conceptual and technical simplicity.\nFor 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\\|\\|$.\nWe hope this explanation is helpful.\n\n> 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?\n\nThank you for raising this valuable question. \nThrough theoretical analysis, we find that the answer is no. \nHowever, 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.\nHere we give the justification:\nLet $\\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}$.\nWe 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$. \nThen 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.\nBy 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)$.\nIn 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.\nMoreover, 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. \n\nIn the simple classic Linear-TD, the same argument holds:\nLet $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.\n\n> 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?\n\nThank 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. \nThus, introducing distribution estimation cannot alleviate these problems.\n\nIn 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.\n\nWe hope this response is helpful.\n\n[1] Marc G Bellemare, Nicolas Le Roux, Pablo Samuel Castro, and Subhodeep Moitra. Distributional reinforcement learning with linear function approximation. AISTATS 2019.\n\n[2] Clare Lyle, Marc G Bellemare, and Pablo Samuel Castro. A comparative analysis of expected and distributional reinforcement learning. AAAI 2019.\n\n[3] Marc G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. MIT Press, 2023. http://www.distributional-rl.org.\n\n[4] Mehrdad Moghimi and Hyejin Ku. Beyond CVar: Leveraging static spectral risk measures for enhanced decision-making in distributional reinforcement learning. ICML 2025.\n\n[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.\n\n[6] Mark Rowland, Marc Bellemare, Will Dabney, R´emi Munos, and Yee Whye Teh. An analysis of categorical distributional reinforcement learning. AISTATS 2018."}, {"review_id": "4Nnop0KjtO", "reviewer": "Reviewer_PXyE", "summary": "This paper studies the sample complexity of temporal difference algorithms based on categorical representations and linear function approximation in the context of distributional reinforcement learning. \nIn particular, the authors analyze the non-asymptotic convergence rate of the proposed method, Linear-CTD, in terms of $W_{1,\\mu_{\\pi}}$ and derive tight sample complexity bounds.", "questions": "Could the author elaborate on the meaning of Equation (10) and the footnote that follows it? \nSpecifically, is the term $(k+1)/(K+1)$ used as a technical convenience for representing the discrete uniform distribution in terms of $\\mathrm{1}_k^T$? \nIf so, does it also help improve the tightness of the derived sample complexity bounds?", "limitations": "Line 214 mentions the need for future analysis of categorical algorithms with nonlinearity.\n\nLine 266 suggests that variance-reduction techniques may help overcome the sample size barrier.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n- Despite the complexity of the literature, the paper clearly explains the necessary background on linear function approximation and distributional RL using standard notation. The overall structure is well-organized and readable.\n\n- The proposed method, Linear-CTD, is robust in the sense that its step size $\\alpha$ does not depend on the number of support points $K$. Although the paper is theory-focused, the authors also provide supporting experiments and include code in the supplementary material.\n\n- By thoroughly comparing with existing analyses on categorical representations, the authors clearly highlight the novelty of Linear-CTD. Although I did not check all proofs in detail, the theoretical results appear consistent with prior work.\n\n**Weaknesses**\n\n- 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.\n\n**Minor Comments**\n\nSince the paper focuses on categorical representations, the authors may consider citing the following two important works in the related literature section:\n\n- Rowland, Mark, et al. \"Statistics and samples in distributional reinforcement learning.\" International Conference on Machine Learning. PMLR, 2019.\n\n- Cho, Taehyun, et al. \"Tractable and Provably Efficient Distributional Reinforcement Learning with General Value Function Approximation.\" arXiv preprint arXiv:2407.21260 (2024).", "was_revised": true, "final_justification": "After reading the rebuttal and participating in the discussion, I am satisfied with the authors’ clarification, particularly regarding the role of normalization and the use of the uniform distribution. I am maintaining my positive score and increasing my confidence from 3 to 4.", "rebuttal": "Thank you for your encouraging comments. \nWe 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.\nRegarding the weaknesses, minor comments and questions, we provide the following detailed responses:\n\n> 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.\n\nThank you for your feedback. We fully agree that adding more explanations for the linear-categorical projected Bellman equation (Theorem 3.1) will enhance readability. \nIn 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.\nIn 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.\nTo further help readers understand, we will add the following table in the revised version that compares various concepts between Linear-TD and Linear-CTD. \nWe believe this will improve the readability of the paper.\n\n| Concepts | Linear-TD | Linear-CTD |\n| --- | --- | --- |\n| Parametrization | $V_\\psi(s)=\\phi(s)^\\top \\psi \\quad [\\text{Eqn.~(3)}]$ | $F_k(s; \\theta)=\\phi(s)^\\top \\theta(k)+\\frac{k+1}{K+1} \\quad [\\text{Eqn.~(10)}]$ |\n| Bellman Operator | $(T^\\pi V)(s)=\\mathbb{E}[r_0+\\gamma V(s_1)\\mid s_0=s]$ | $(\\mathcal{T}^\\pi \\eta)(s)=\\mathbb{E}[(b_{r_0,\\gamma})\\_{\\\\#}\\eta(s_1)\\mid s_0=s]$ |\n| Projection Operator | $V_{\\tilde \\psi}=\\Pi^\\pi\\_\\phi V, \\tilde{\\psi}=\\Sigma^{-1}\\_{\\phi}\\mathbb{E}_{s}[\\phi(s)V(s)]$ | $\\eta_{\\tilde{\\theta}}=\\Pi^{\\pi}\\_K \\eta, \\tilde{\\Theta}=\\Sigma^{-1}\\_{\\phi}\\mathbb{E}_{s}[\\phi(s)(F\\_{\\eta}(s)-F\\_{\\nu}(s))^\\top]  \\quad [\\text{Eqn. (11)}]$ |\n| Projected Bellman Equation | $V_{\\psi}=\\Pi^\\pi_\\phi T^\\pi V_{\\psi} , \\quad [\\text{See Eqn.~(4)}]$ | $\\eta_{ \\theta}=\\Pi^\\pi_{\\phi,K} \\mathcal{T}^\\pi \\eta_{\\theta}, \\quad [\\text{See Eqn.~(12)}]$ |\n| Update Rule | $[\\text{See Eqn.~(5)}]$ | $[\\text{See Eqn.~(13)}]$ |\n| $A_t$ | $\\phi(s_t)\\phi(s_t)-\\gamma\\phi(s_t)\\phi(s_{t+1})$ | $[I_K\\otimes(\\phi(s_t)\\phi(s_t))]-[(C\\tilde G^{-1}(r_t)C^{-1})\\otimes(\\phi(s_t)\\phi(s_{t+1}))]$ |\n| Key Quantity in $A_t$ | $\\gamma$ | $C\\tilde G^{-1}(r_t)C^{-1} \\text{ with spectral norm }\\sqrt{\\gamma}$ |\n| $b_t$ | $r_t\\phi(s_t)$ | $\\frac{1}{K+1}[C(\\sum_{j=0}^Kg_j(r_t)-1_K)]\\otimes \\phi(s_t)$ |\n| Key Quantity in $b_t$ | $r_t\\leq 1$ | $K^{-3/2}(1-\\gamma)^{-1}\\Vert C(\\sum_{j=0}^Kg_j(r_t)-1_K)\\Vert \\leq 1$ |\n| Measure of Error | $\\Vert V-V^\\pi\\Vert \\_{\\mu\\_\\pi}=(\\mathbb{E}_{s}[(V(s)-V^\\pi(s))^2])^{1/2}$ | $W\\_{1,\\mu_\\pi}(\\eta, \\eta^\\pi)=(\\mathbb{E}_{s}[W_1^2(\\eta(s),\\eta^\\pi(s))])^{1/2}$ |\n| Approximation Error | $\\Vert V\\_{\\psi^\\star}-V^\\pi\\Vert \\_{\\mu_\\pi}\\leq (1-\\gamma^2)^{-1/2}\\Vert \\Pi^\\pi\\_\\phi V^\\pi-V^\\pi\\Vert \\_{\\mu_\\pi}$ | $[\\text{See Proposition~3.2}]$ |\n| Sample Complexity | $\\widetilde{\\mathcal{O}}\\left(\\frac{\\Vert \\psi^\\star\\Vert ^2_{\\Sigma_{\\phi}}+1}{(1-\\gamma)^2\\lambda_{\\min}}\\left(\\frac{1}{\\varepsilon^2}+\\frac{1}{\\lambda_{\\min}}\\right)\\right)$ | $\\widetilde{\\mathcal{O}}\\left(\\frac{\\Vert\\theta^\\star\\Vert ^2_{V_1}+1}{(1-\\gamma)^2\\lambda_{\\min}}\\left(\\frac{1}{\\varepsilon^2}+\\frac{1}{\\lambda_{\\min}}\\right)\\right)\\quad [\\text{See Theorem~4.1}]$ |\n\n> Minor Comments\n\nThank you for your reminder. \nThe 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. \nWe will discuss these works in the revised version.\n\n> 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?\n\nThank you for your valuable questions and suggestions. \nThe 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$.\nAnd as you noted, introducing normalization indeed helps improve the tightness of the derived sample complexity bound and is indispensable. \nIn addition, the normalization also offers algorithmic and interpretability benefits. \nWe will elaborate on normalization in detail below and add this content in the revised version.\nWe believe that this will enhance readers' understanding.\n\nAs 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$.\nFirst, we explain why normalization is necessary.\nAlthough we work with signed measures, it remains necessary to ensure the total mass is normalized to $1$.\nThe reasons are from both algorithmic and theoretical considerations.\nFrom an algorithmic perspective, without normalization, the difference between our estimator and the true return distribution would not be a signed measure with mass zero. \nThis would make our estimator less interpretable and lead to issues in defining many important quantities. \nFor 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$. \nFurthermore, 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.\nIn this sense, the algorithm will not be a natural extension of Linear-TD and tabular CTD.\nFrom a theoretical perspective, tight sample complexity bounds can only be achieved after introducing normalization. \nWithout 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.\n\nNext, we explain the impact of choosing different normalizers (distributions used for normalization) and why we choose the discrete uniform distribution $\\nu$ as the normalizer. \nAs 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. \nFortunately, such undesirable properties would not appear in the tabular case.\nIn 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. \nIn 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). \nBased on the above discussion, any distribution could be used for normalization in practice.\n\nAnd we choose the uniform distribution as the normalizer only for conceptual and technical simplicity.\nFor 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\\|\\|$.\nWe hope this explanation is helpful.\n\n[1] Rowland, Mark, et al. Statistics and samples in distributional reinforcement learning. International Conference on Machine Learning. PMLR, 2019.\n\n[2] Cho, Taehyun, et al. Tractable and Provably Efficient Distributional Reinforcement Learning with General Value Function Approximation. arXiv preprint arXiv:2407.21260 (2024)."}, {"review_id": "NVlVisOTIH", "reviewer": "Reviewer_j4KH", "summary": "The paper investigates distributional temporal difference learning in the presence of linear-function approximation and proposes the Linear-CTD algorithm, which enjoys the same sample complexity as classical TD, albeit estimating the entire return distribution as opposed to just its expectation. This is done by viewing the categorical Bellman equation as a linear system and by preconditioning the matrix  so that convergence depends no longer on the number of support points  K; rigorous non-asymptotic high-probability results are further obtained by studying exponential stability of random matrix products and applying Polyak-Ruppert tail averaging, with empirical researchers confirming these rates and showcasing gains in performance over previous methods.", "questions": "see weakness part", "limitations": "see weakness part", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths：1. the first non-asymptotic high-probability bounds for distributional TD with linear function approximation, thus filling in the gaps  2. both the linearization and the analysis of exponential stability have been carried out with great care; also, the bounds are consistent with those known in the Linear-TD setting.\n\nWeaknesses：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. \n\n2. However, there is no compelling need shown in this paper on why an in-depth analysis of distributional convergence is necessary.", "was_revised": true, "final_justification": "Most of my concerns are addressed.", "rebuttal": "We sincerely appreciate your effort and time in reviewing our paper.\nAnd we are glad to hear that you find our results are novel and technically sound. \nBelow, we hope our responses can address your concerns.\n\n> 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.\n\nWhile 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.\n\n1) 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]. \nTo 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).\n\n2) 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. \nThese 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.\nWe believe these results are of independent interest and will be beneficial for subsequent research in distributional reinforcement learning.\n\nWe also would like to mention our algorithmic contributions, which is fully original.\nTo be specific, we proposed an improved version of the linear-categorical TD learning algorithm (Linear-CTD). \nCompared 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.\n\nIn summary, we believe that our work demonstrates significant originality and contributions compared to pre-existing results. \nWe hope that the response clarifies this point.\n\n> Weakness 2. However, there is no compelling need shown in this paper on why an in-depth analysis of distributional convergence is necessary.\n\nWe would believe in-depth theoretical analyses of distribution policy evaluation algorithms is timely and important.\nIn 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]. \nBuilding 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. \n\nIn 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.\n\n1. 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). \nDistributional convergence results offer theoretical guarantees for the reliable estimation of these risk metrics in practical tasks.\n\n2. Many risk-sensitive RL algorithms [13, 14] incorporate distributional RL as a core step to estimate specific expected utility functions of return distributions.\nThe distributional convergence results can provide finite-sample convergence rates for this critical step, giving theoretical guarantees for these algorithms.\n\nFor 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].\n\n[1] Marc G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. MIT Press, 2023. http://www.distributional-rl.org.\n\n[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.\n\n[3] Gen Li, Changxiao Cai, Yuxin Chen, Yuting Wei, and Yuejie Chi. Is q-learning minimax optimal? a tight sample complexity analysis. Operations Research,\n72(1):222–236, 2024a.\n\n[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.\n\n[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.\n\n[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.\n\n[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.\n\n[8] Liangyu Zhang, Yang Peng, Jiadong Liang, Wenhao Yang, and Zhihua Zhang.\nEstimation and Inference in Distributional Reinforcement Learning. The\nAnnals of Statistics, 2025. Forthcoming.\n\n[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\n\n[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.\n\n[11] Haoran Wang and Xun Yu Zhou. Continuous-time mean–variance portfolio selection: A reinforcement learning framework. Mathematical Finance, 30(4): 1273–1308, 2020.\n\n[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.\n\n[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.\n\n[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.\n\n[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."}, {"review_id": "f0q3AK8NFl", "reviewer": "Reviewer_VGyk", "summary": "This paper investigates the finite-sample statistical rate of distributional TD Learning under linear function approximation. The authors propose a Linear-CTD which achieves finite-sample convergence rates that are independent of the number of categorical supports, a key improvement over vanilla approaches. They establish the first non-asymptotic sample complexity bounds for distributional TD and derive finite-sample rates that match those of classical TD learning, demonstrating that estimating the full return distribution is statistically no hard than estimating its expectation. Theoretical contributions are complemented by empirical results validating the improved performance and robustness of Linear-CTD.", "questions": "N/A", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "In my opinion, the novelty of this paper lies in its theoretical development, which establishes the first non-asymptotic sample complexity bounds for distributional TD learning under linear function approximation. This fills the gap in the RL. The authors show that the results are consistent with classical TD rates, providing the important insight that learning the full return distribution is no harder than learning its expectation. The paper is well written and clearly organized. In particular, I like Section 5, which outlines the proof strategy, making the non-asymptotic analysis transparent and easy to follow. Theoretical findings are also supported by empirical experiments, which reinforce the effectiveness of the proposed approach. I do not see major weaknesses in the paper, although I may overestimate the novelty and miss some technical issues, as I did not check the proofs line by line.", "was_revised": true, "final_justification": "I have no remaining questions and would keep my score.", "rebuttal": "Thank you for your encouraging comments and constructive suggestions.\nWe 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. \nWe are also grateful for your positive feedback on our writing (particularly Section 5) and the experimental section.\n\nRegarding your concerns, we provide the following responses:\n\n> 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).\n\nThanks 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.\n\nWe will add the following table in the revised version.\nIn 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.\nThis allows for a clearer comparison of theoretical results with prior work.\n\n| Paper | Sample Complexity | Method |\n| --- | --- | --- |\n| [1, 2, 3] | $\\widetilde{\\mathcal{O}}\\left(\\frac{\\|\\|\\psi^\\star\\|\\|^2\\_{\\Sigma_{\\phi}}+1}{(1-\\gamma)^2\\lambda\\_{\\min}}(\\frac{1}{\\varepsilon^2}+\\frac{1}{\\lambda\\_{\\min}})\\right)$ | Linear-TD |\n| [4] | $\\widetilde{\\mathcal{O}}\\left(\\frac{\\|\\|\\psi^\\star\\|\\|^2\\_{\\Sigma\\_{\\phi}}+1}{\\varepsilon^2(1-\\gamma)^2\\lambda\\_{\\min}}\\right)$ | Linear-TD with Variance Reduction |\n| [5] (Section 9.7) | Contraction Analysis | Vanilla Linear-CTD |\n| This Work (Theorem E.2) | $\\widetilde{\\mathcal{O}}\\left(\\frac{K^4(\\|\\|\\theta^\\star\\|\\|^2\\_{V_1}+1)}{(1-\\gamma)^2\\lambda\\_{\\min}}(\\frac{1}{\\varepsilon^2}+\\frac{K^2}{\\lambda\\_{\\min}})\\right)$ | Vanilla Linear-CTD |\n| This Work (Theorem 4.1) | $\\widetilde{\\mathcal{O}}\\left(\\frac{\\|\\|\\theta^\\star\\|\\|^2\\_{V_1}+1}{(1-\\gamma)^2\\lambda\\_{\\min}}(\\frac{1}{\\varepsilon^2}+\\frac{1}{\\lambda\\_{\\min}})\\right)$ | Linear-CTD |\n\nHere, 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). \nThe contraction analysis in [5] in the table means that they provided a contraction analysis for the dynamic programming version of Vanilla Linear-CTD algorithm.\n\n> 2. Additionally, the authors may consider presenting some of the empirical results in the main text, despite space limitations.\n\nThank you for your suggestion. We agree that presenting some experimental results in the main text would be helpful for readers. \nWe will move Figures 1 and 2 to the main text in the revised version.\nFigure 1 demonstrates the convergence results of our algorithm. \nBy 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). \nIn contrast, the step size of our Linear-CTD does not require adjustment as $K$ increases (Lemma 5.3).\n\n[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.\n\n[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.\n\n[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.\n\n[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.\n\n[5] Marc G. Bellemare, Will Dabney, and Mark Rowland. Distributional Reinforcement Learning. MIT Press, 2023. http://www.distributional-rl.org."}]

Author rebuttal:

Meta-review: This paper analyzed the finite sample convergence of distributional TD learning with linear function approximation. A particular useful take-away is that distributional TD is not statistically harder (ie., taking more samples) than standard TD, thus allowing bootstrapping more information for applications where higher order statistics are useful for decision making.

The reviewers appreciate the clear technical contribution of the paper and the reviewer-author discussion has resolved most of the concerns. Therefore, the paper is worthy of acceptance.
Decision: Accept (poster)