Categorical Distributional Reinforcement Learning with Kullback-Leibler Divergence: Convergence and Asymptotics
Authors: Tyler Kastner
Mark Rowland
Yunhao Tang
Murat A Erdogdu
Amir-massoud Farahmand
Venue: ICML 2025 poster (2025)
Keywords: Distributional Reinforcement Learning
TD learning
Categorical Temporal Difference learning
Primary area: reinforcement_learning
TL;DR: We study categorical distributional reinforcement learning with a KL divergence-based loss, providing theoretical guarantees, empirical validation, and actionable insights for practitioners.

Abstract: We study the problem of distributional reinforcement learning using categorical parametrisations and a KL divergence loss. Previous work analyzing categorical distributional RL has done so using a Cramér distance-based loss, simplifying the analysis but creating a theory-practice gap. We introduce a preconditioned version of the algorithm, and prove that it is guaranteed to converge. We further derive the asymptotic variance of the categorical estimates under different learning rate regimes, and compare to that of classical reinforcement learning. We finally empirically validate our theoretical results and perform an empirical investigation into the relative strengths of using KL losses, and derive a number of actionable insights for practitioners.

Categorical Distributional Reinforcement Learning with Kullback-Leibler Divergence: Convergence and Asymptotics
Tyler Kastner 12 Mark Rowland
Yunhao Tang*1
Murat A. Erdogdu 12 Amir-massoud Farahmand 561
Abstract
ing methods cast the return distribution prediction problem as one of classification. These methods have proven perfor- mant in a variety of large-scale settings.
We study the problem of distributional reinforce- ment learning using categorical parametrisations and KL divergence loss. Previous work ana- lyzing categorical distributional RL has done so using Cramér distance-based loss, simplifying the analysis but creating a theory-practice gap. We introduce a preconditioned version of the al- gorithm, and prove that it is guaranteed to con- verge. We further derive the asymptotic variance of the categorical estimates under differentt learn- ing rate regimes, and compare to that of classical reinforcement learning. We finally empirically validate our theoretical results and perform an em- pirical investigation into the relative strengths of using KL losses, and derive a number of action- able insights for practitioners.
Theoretical convergence analysis has also been established for categorical dynamic programming, and categorical temporal-differeee learning using the Cramér loss (Row- land et al., 2018; Boeck & Heitzinger, 2022; Peng et al., 2024). However, most large-scale implementations of cate- gorical temporal-difference difference learning use a KL loss, rather than Cramér loss. This is a crucial detail of large-scale im- plementations, but has not yet been theoretically analysed.
In this work, we study categorical temporal-difference learn- ing with KL loss (KL-CTD) as a fundamental tabular al- gorithm for reinforcement learning in its own right. In Section 3, we present several empirical examples of intrigu- ing behaviour of KL-CTD in comparison to classical TD learning, motivating our study. We then go on in Section 4 to establish a connection between KL-CTD and distribu- tional dynamic programming algorithms, which allow us to characterise the long-term behaviour of KL-CTD condi- tional on its convergence. In Section 5, study the ques- tion of convergence, and introduce novel preconditioned variant of KL-CTD. We then provide theory on the finer- grained asymptotic fluctuations of value function estimates produced by KL-based categorical distributional algorithms in Section 6. Our results provide a variety of hypotheses about the behaviour of KL-CTD, when we should expect it to perform well relative to classical TD, and intuition on how to set hyperparameters such learning rate. We then validate these hypotheses in variety of experiments.
1. Introduction
The expected return is core object in reinforcement learn- ing, allowing for both the evaluation of agent's 's behaviour, and providing means of improving this behaviour (Sutton, 2018). The traditional approach to predicting the expected return is to directly model it as mcan- peean-predictin regression problem. Distributional reinforcement learning algorithms take a different approach, instead predicting the full proba- bility distribution of the random return; the expected return then emerges as byproduct of this richer prediction. Agents making use of distributional reinforeement learning have enjoyed variety of empirical successes (Bellemare et al., 2017; Hessel et al., 2018; Yang et al., 2019; Bodnar et al., 2020; Wurman et al., 2022).
2. Background
We consider a Markov decision process (Puterman, 2014; Sutton, 2018) with finite state space X' action space A. transition matrix P E XxXXXX reward kernel R X' x A + ( and discount factor 7. We as- sume a fixed policy : X + (AA,, which gives rise to the policy transition matrix p= E xxxx (where pa (z'|z) = EEEA P(11x) a)) and the policy re- ward kernel R X - (R) (where R² = EneA ((az))((2)))
In particular, categorical distributional reinforcement learn-
'University of Toronto VVector Institute 'Google DeepMind "Meta Platforms, Inc. *Work done while at Google Deep- mind 'Polytechnique Montréal "Mila. Correspondence Tyler Kastner <tkastner@ toronto.cduz, Mark Rowland <markrow- land www.ggogoogg..m...
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).
The policy n induces a distri-

---

Categorical Distributional RL with KL Divergence
bution over trajectories (X,, Ac, Re)120- where for every t 2 0 have A, x( Xt), R R(X,, At), and Xx+1 P(- P(, The agent's performance is summa- rized by the discounted return E(2o?' RR. We use G" (x: to denote the random return when generating actions accord- ing to т and beginning the trajectory in state Xu = 2, and write 1f" (x) for the distribution of this random variable. The expected return across trajectories is the value function
where 21,...: 2m E R are fixed locations, and ((p(())) 2: E X)) are learnable probability parameters. Throughout this paper, we will consider the case where each vector (p)(z)) is parametrized via a vector of logits 6(x) E I"", writing (x)) = softmax(6(x)) (Bellemare et al., 2017).
Categorical MC estimation. The categorical counterpart to the Monte Carlo update in Equation (1) is based on the idea of regressing the distribution n7()) = > (( towards the empirically observed ōg. In practice, we typ- ically have G g {= Am J. and so the KL divergence between åg and n(x) is infinite. To circumvent this issue, the outcome G is first mapped to a two-hot distribution as illustrated in Figure/ (Bellemare et al., 2017; Schrittwieser et al., 2020), producing the projected target distribution
A (z) = E G" (x)]
estimating this quantity is central to value-based RL.
2.1. Value estimation
The Monte Carlo (MC) approach to estimating V performs online regression against sample returns from trajectories. More precisely, we maintain an estimate V F of the value function VT, and given return G (x) generated by following π from an initial state z, we perform the update
E h2 (
(4)
The precise definition of the (ht))= functions defining the probability masses in Equation (4) is given by
V(+) + V(+) a (G Vxx))
(1)
2 24-1
2i+1 - 2
h₂(2) = mmxx 0, min
where a > 0 is step size parameter.
2i 2i-1
Zi+1 - 24
Temporal-diffeen difference (TD) methods utilize bootstrapping in place of sample returns from completed trajectories. More
for 2 Sis m 1, and special edge cases
precisely, having observed sample transition RR, X') generated by following the policy π,, the full sample of the return G in the MC update of Equation (1) is replaced with a bootstrap estimate + (X'), the intuition being that if V x V*, then Ei + /V (X') s E[G]. This leads to the TD update equation
h1 h1(2) min
1, 1,
0, 2/
22 - 21
2 2m 1
hm min
max
0,
2m Zm
which account for "clipping" outcomes that occur outside of the interval [21, 2m] to the endpoints.
V(x) + V(z) a
a (R+ (R+>V()))))) V(z)).
learning can deliver several benefits rel- ative to MC estimation, such as reduced variance in regres- sion targets, as well the ability to learn from individual transitions rather than requiring full trajectories.
(G)
2,1
Zi+1
mm
G
2.2. Categorical distributional reinforcement learning
Figure 1. Projection of the outcome G onto a distribution over the support set "mm
As opposed to classical value-based RL algorithms that focus solely estimating the expected return of pol- icy, distributional reinforcement learning algorithms learn parametrized approximation of the return distributions, i.e., (x) Dist ( (r)),, for each æ E X.
The categorical Monte Carlo (CMC) update, in analogy with the Monte Carlo update in Equation (1), can then be defined as
Various families of distributional RL algorithms have been proposed, making use of distinct parametric families of re- turn distributions (Morimura et al., 2010b; Dabney et al., 2018). A popular class of methods cal parametrizations (Bellemare et al., 2017), in which the distributions (17* (x) x E X)) are approximated through categorical distributions of the form
Φ(z) - d(x) V VaKL(h(G) (l
(5)
where write h(G) = (h, ())) for the vector of two- hot probabilities corresponding to the outcome G.
based categori-
Categorical TD learning. Just as in the case of value estimation described in Section 2.1, the CMC update can modified to allow for learning from individual sample transitions (I,, R, X'). As with classical TD learning, the
πx)) = Σm (2
2

---

Categorical Distributional RL with KL Divergence
sampled return G is replaced by bootstrapped equivalent. Here, the random return G is approximated by R yZ', with Z' sampled from the categorical approximation r7(') to the return distribution at state X'. Setting up the KL gradient update as in Equation (5) for these bootstrapped samples, averaging according to the probability of sampling Z' = 2j for: I,...: m, then yields the following update:
10
MSE
10
TD
Cramér-CTD
KL-CTD
10a
10
102
10a
m
((x) ((x) Σ
P (x')x
(6)
Iteration
Figure 2. Mean-squaree error of TD, Cramér-CTi and KL-CTD with tuned learning rates.
VoKT,(uR+ + p (x))
This is the categorical temporal-din -enporal-difference (CTD) learning update, proposed by Bellemare et al. (2017). This update forms& core component of many deep reinforcement leam- ing agents utillizing distributional RL (Bellemare et al., 2017; Hessel et al., 2018). Our principal goal in this paper is to consider categorical cunporal-different learning as a fun- damental, complementary approach value prediction, and to understand its relative strengths and weaknesses pared to classical TD learning.
VT
1.5
1.0
The Cramér-CTD) update. While the CTD update in Equa- tion (6) has not previously been analysed, Rowland et al. (2018) prove convergence of an alternative version of cate- gorical temporal-diffferece learning that performs a mixture update on probabilities directly, given by:
0.0
-0.2
0.0
0.2
0.4
0.6
0.8
V(x1)
p(x) - p(x) + a (X")(R +22)) 724) - -p(x))
Figure 3. Value function dynamics of KL-CTD 2-state MDP.
We refer to this variant as Cramér-CTD. CTD, due to the fact that it can be derived through the use of the Cramér distance (Cramér), rather than the KL, as a loss, and refer the update in Equation (6) KL-CTD, in contrast.
Another striking difference of KL-CTD is the behaviour of its mean estimates. In particular, as a result of the loss being optimized is the KI. KI. divergence, which generally does not correspond to an improvement in the mean estimate. For probability measures P. Q supported on [a, b] with means (Pp; lQQ we always have that
3. Motivation for the study of KL-CTD
Before beginning our analysis, we pause to motivate the study of KL-CTD, despite existing work studying Cramér- CTD (Rowland al., 2018) and proving its equivalence to TD-learning in the tabular setting (Lyle et al., 2019). We will show unlike Cramér-CTD, KL-CTD has notable algo- rithmic differences to TD-learning, and understanding these differences is consequential as KL-CTD is the foundation of many performant deep RI. algorithms (Bellemare et al., 2017; 7; Hessel et al., 2018; Farebrother et al., 2024).
IPP uQl S (b 10-aay (XT(2K](PP))))
so that optimizing the KL divergence will eventually op- timize the mean estimate, however this is not guaranteed at each step. In the Monte Carlo setting, in expectation and under a suitable learning rate, the value eтror of the MC estimate is guaranteed to decrease at each step (see Appendix A). In contrast, we find the value error of the categorical Monte Carlo estimates may behave erratically as seen in Figure 3, entirely due to the dynamics optimizing the KL. divergence.
We present an example of learning curves for TD, Cramér- CTD, and KL-CTD on a particularly-chosen environment (see Appendix G) in Figure? 2. Learning rates were tuned independently for each algorithm. The curves for TD and Cramér-CTD are perfectly overlapped, as predicted by Lyle et al. (2019). Further, KL-CTD exhibited significantly dif- ferent learning dynamics, achieving a much better value MSE on this particular MDP.
4. A linear-algebrait perspective on KL-CTD
We begin by developing deeper understanding of the KL- CTD update rule in Equation (6), and establishing connec- tions with existing dynamic programming algorithms.
3

---

Categorical Distributional RL with KL Divergence
Proposition 4.1. The expected KL-CTD update can be ex- pressed via matrix-vector multiplication
Interestingly, the form of Equation (7) shows that the (ex- pected) KL-CTD updates are linear in probability space, and the nonlinearity of the updates is entirely due to the softmax transformation of the logits. We believe that this perspective may be interesting in its own right, as this form of nonlinearity is quite distinct from other nonlinear distribu- tional al algorithms, such as quantile-based As a result of this connection between KL-CTD and dynamic programming. conclude the following result.
c + a(T - -)pp,
where we interpret & and pd as [X] amm-immension. vectors, and where T is the categorical distributional Bellman operator (Rowland al., 2018; Bellemare et al., 2023; Rowland et al., 2024), given by
Proposition 4.3. The only stationary point of expected up- date in Equation (7) is the fixed point 8 of...
T (1,, v. PT (x) (B) 7-
The connection that Proposition 4.1 establishes the cate- gorical distributional Bellman operator is important because a rich convergence theory exists for this operator (Rowland et al., 2018). Intuitively, repeated application of the opera- tor 7 to any initial collection of categorical probabilities p E IR xm results in convergence to collection of proba- bilities p" such that the associated categorical distributions
Therefore, if KL-CTD converges to stationary point, it must converge to D", and it thus produces accurate esti- mates of return distributions (with error as stated in Propo- sition 4.2), and the support (=I,..., 2m } is correctly se- lected to include all possible returns, then its mean predic- tions at convergence are exact. This our first important finding: although the transient mean dynamies of KL-CTD may deviate significantly from those of classical TD (as shown in Section 3), at convergence it produces exact pre- dictions, in contrast to other distributional approaches such as quantile TD (Dabney et al., 2018; Rowland et al., 2023).
- m P. (3)
form close approximations to the true return distributions 77" (ar). We recall the precise results required in the remain- der of the paper below.
5. Asymptotic convergence and preconditioned KL-CTD
Proposition 4.2 (Rowland et al., 2018). The mapping T A + A is contraction with respect to the norm I|
defined by
We now turn our attention to questions of convergence for categorical algorithms using KL. divergence.
P = max [Cp(z))l2, ZEA
(8)
5.1. Categorical Monte Carlo
where |1 is the standard Euclidean norm, and C is the lower-triangular matrix with Cij i 2 j, and Cij otherwise. More precisely. for any p.q E A X
The categorical MC update is straightforwardly analysable, as the following proposition shows.
m. have
Proposition 5.1. IfG v, then the CMC update appearing in Equation (5) is stochastic gradient for the objective KL(E@wu [h[G) || pm), which is convex in Ø.
p T" --1lp-10 4|1-
As a result, T has a unique fixed point p" E A Writing os"t E P (B.xx for the corresponding distributions, so that
Stochastic gradient descent on convex objectives is well- behaved theoretically (Bottou et al., 2018), and under mild conditions, it follows that iterative CMC updates with inde- pendent samples results in pd converging.
= E.t (
we have that if [21, 2m] contains all possible returns in the environment, 77" (x) has the same mean as 17" (x), and moreover,
5.2. KL-CTD and preconditioning
Next, consider the KL-CTD algorithm. In Section 2, we motivated this algorithm variant of CMC that introduces bootstrapping, a central idea in RL where the algorithm's 's own predictions are used as targets. Bootstrapping means that the algorithm updates can no longer be interpreted as stochastic gradient descent, and the question of whether such an algorithm converges therefore more subtle.
(T), (z)) (m ()m-1)(1-)) 1))
where l2 is the Cramér distance (Cramér; Székely, 2003; Székely & Rizzo, 2013), defined by
1/2
62(u,1') = (1000))--
A heuristic argument for preconditioning. A natural question is whether the expected updates of KL-CTD still act to reduce quantity such as Σ KL. (x) (x));
and Fu, Fut are the CDF' of v, v', respectively.
4

---

Categorical Distributional RL with KL Divergence
such observation would form the basis of a proof of convergence for KL-CTD. However, this does not in fact
11 is the norm defined by l|lll Exex dt (x)][C((x>]
where
hold in general; we give a numerical example in Appendix A (although we emphasize that this doesn't t imply that KL- CTD does not converge). The issue stems from the fact that the expected change this quantity
Next, we establish that weighted sum of KLs between the CDP fixed-point p" and the approximation pd must strictly decrease when pd is updated with the continuous- time dynamics associated with Equation (10).
Σ TEN
(=)).
Proposition 5.4. The continuous- dynamics of Equa- rion (10) converges, and a Lyapunov function for its conver- gence is
=a(p - p".1"pp pp)
pd) O(a))
In general, the inner product appearing above may be pos- itive, leading to an increase in the summed KLs. This is due to the fact that T may be strictly expansive in the ? norm on return PMFs. However, the background theory in Proposition 4.2 shows that 7 exhibits contractive be- haviour in a weighted norm. This motivates preconditioning the KL-CTD update with matrix that incorporates this weighting factor, so that expected updates under this new preconditioned rule reduce KL-like measurement to p" by construction. We introduce this new algorithm below.
L( = (
(z)
(1I)
ZER'
where dt is the stationary distribution of the policy и.
With these results, we are now ready to state the central con- vergence theorem for PKL-CTD, in the case synchronous updates. We discuss the generalization to the asynchronous updates setting in Appendix A.3.
Definition 5.2. Given logits & E IR parametrising cat- egorical approximations of return distributions as in Equa- tion (3), and an observed transition (x, R, X'), the precon- ditioned KL categorical ceppora--differene learning (PKL- CTD) update is defined by
Theorem 5.5. Suppose that (œx))20 is a sequence of logits generated according to synchronous PKL-CTD updates.
That is, for each ks 2 0, have independent transitions (x, Rt, Xg) such that
Φk+1(T) =6((T) axx m C tl (XE) + + -p - -x)
((x) + 6(x)
E P (x))x
(9)
Σ
CL
(x)).
Further suppose that the stepsizes (ou))20 satisfy the Robbins-Munro conditions _k= an = oo and Ex=o a < Then we have that Ob converges in the sense that p pr (x) for every x E X almost surely.
We explore the empirical performance of this update in Section 7: in the remainder of this section, we now inves- tigate its convergence properties, describing how to make the motivation intuition above precise. Our proof of conver- gence follows the Slobins-Siegmuna theorem (Robbins & Siegmund, 1971) 1) and a chaining argument (Thakoor et al., 2022); to understand the dynamics of the random updates in Equation (9), we first analyse the dynamics of the corre- sponding expected updates. By following Proposition 4.1, we obtain an expression for the expected PKL-CTD update in terms of the operator T"
6. Analysis of value estimates
We now perform a theoretical investigation into the algo- rithms considered and introduced, and their efficacy as value estimates. We consider both Monte Carlo (Section 6.1) and TD (Section 6.2) settings, and these analyses yield theoretically-motiiatee insights for practitioners.
(x+1(x) = (x(x) + CCC - -1)p²)(x)
(10)
6.1. Monte Carlo setting
We then establish two results that are key in establishing the convergence of PKI-CTD. The first establishes a new weighted contraction property for 7",, mirroring analysis undertaken by Wu et al. (2023) in the case of I-Wasserstein distance.
We begin by examining the Monte Carlo methods. For the re- sults in this section, we will write (Vx)a20 and (GCMC))s) for sequences of iterates generated by following Equa- tion (1), Equation (5), respectively. Further, we will con- sider a Monte Carlo variant of PKL-CTD (preconditioned CMC; PCMC), defined by replacing the bootstrap samples in Equation (9) with Monte Carlo sample G, yielding the update
Proposition 5.3. The mapping T A A is a con- any
traction with respect to the norm || ||x:: for any p. q AX we have
1T - T gll 5 171pp-1000000000 g||=,
6(x) + ddx)) C
(x)). (12)
5

---

Categorical Distributional RL with KL Divergence
We write (oPCM))
120 for the sequence of iterates gener- ated by following Equation (12). The stationary points for both CMC and PCMC are any logits such that the induced probabilities are equal to Io (x) = B (x)) [h(G) at each state, which we will refer to as 1j" (x).
We next consider the second differentiator, that is the differ- ence of the variances due the Cramér projection.
Proposition 6.3. Suppose is a probability measure whose support is contained in [a, bl, and Varz~u(Z) = o². Let
a = : 7m = be m equally spaced points, Then we have
The categorical estimates are used to produce value esti- mates given by
Varz~Beu o =
VCMC (I)) = = pb CMC
(I), PCMC (x) = = =Tpot (x))
where 0 S E(V) S
1(m-1yz is a quantity capturing the
(b-a)²
P2
amount of projection in the map v + Nev.
where 2 = (2) E T is the vector of locations.
Proposition 6.3 indicates that the projection increases the variance incurred, but the amount of this increase is upper bounded by (b-o)² (m-1)2 This suggests that both (i) choosing the interval [a, 6) as small as possible while still containing the all possible returns (so as to not bias the mean estimate), and (ii) increasing the number of locations will both lower the asymptotic variance of CMC.
The following proposition establishes convergence proper- ties for the decreasing step size case; the analysis builds core results asymptotic fluctuations in stochastic approx- imation theory (Borkar, 2008).
Proposition 6.1. Suppose the iterates (VK)&200 (oR (OCM)) 1220. )x20, and (oPCMC) k20 were produced using step size Og-- = agk- for B (1/2,1, 1). Further suppose that V. (x) > O for each x, E X x [mm] Then we have that
On the number of locations and learning rate. We now consider the constant step size setting, and identify another phenomenon which has important implications for practi- tioners: the role of learning rate and number locations are jointly related in the learning process, and should be jointly tuned. Concretely, (both MC and TD) in the form
1 (x) (x)) 4 NO
(1 (z) (x)) 4 V (0, ]= ((T"
1/0/2(V (VPCMC (x) VT 4N (0. (ux))()) ((()))
0k
can write KI. loss-based updates
where (p) = diag(p) - pp' u(x)= = CJ (x))= and o² the variance of (z))
a - Φ a DV
(13)
The technical assumption on positivity of p7 (x) avoids the case of divergence of logits to too. For both CMC and PCMC, the form of the limiting variance may appear opaque. However, we can note that the variance of the projected return distribution pd (a:) is given by
where D is a general preconditioner and T is a target dis- tribution. The eventual mean estimate is then produced by (po, =). The map & + Pe is vm -Lipschhitz (see Proposi- tion A.1), meaning that an update of order on é results in a change in the mean estimate of order vmm Additionally, the gradient appearing in Equation (13) 7m -Lisschhh so that as the number of locations increases, the loss land- scape becomes proportionally flatter and larger step sizes can be afforded. Based on these points, we suggest seal- ing a by /m, as our analysis suggests that this allows for transferable learning dynamics across different values of m.
Σ pi22
Σ Pizz)
2 = 2
- (=*))=)
= =2T J(p))z,
where we write p; = =pp (x)); ) for clarity. Hence the limit- ing variance of categorical Monte Carlo at a state x differs from the variance of f (x) by multiplication of [ (x)) in the quadratic form. We can now isolate two sources of dif- ference between the limiting variance of CMC and MC: the difference due to the extra multiplication of [ix (z:),, and the difference between the variance of the return distribution 7 (ar) and the variance of the projected return distribution p (z). We begin by analyzing the former, and demonstrate that the extra factor of J [Dx (z:) produces a quantity which is strictly smaller than the variance, yet still continuous with respect to the variance.
6.2. Temporal-ditffrenn setting
We next turn to temporal-di difference learning algorithms. We will write (Ve)a20 for sequence of value estimates following Equation (2) and (oPCTD) for sequence of logits following Equation (9), with associated value esti- mates VPCTD The following results allow us to compare 6e. exact asymptotic variances for TD and PKL-CTD; the result for TD is due to Wu et al. (2024), while the PKL-CTD result is our own.
Proposition 6.2. Suppose p- (pi)m is collection of probabilities associated to the locations = (=1) and let the variance of this categorical distribution be a². Then there exists 3 depending only on p such that
Proposition 6.4. (Wu al., 2024) If the step sizes (ok))20 are given by as agk- -8 we have that
3o² < T J())² < <²
VT (V²)
$ NOO, (ADAEDAA )o),
6

---

Categorical Distributional RL with KL Divergence
where we use the notation A = - pp and ETD is the X) where the expectation is over sample transitions (x, Ra, x²))
N 102
N 10
matrix Plzz"], == (Ry (Rx+
0.30
Thyoratical Density
0.30
Theoratical Density
Empinical Dersity
Empirical Density
0.25
0.25
Proposition 6.5. If the step sizes (ox))20 are given by Ok -ook-,, -, we have that
0.20
0.20
(V (x) v= 3 N (0, by (2)
0.10
0.30-
0.05
0.05
where by J(p (z))z and EPCTD( ) is the unique solu-
0.00
0.00
-2.5
2.5
tion E of the Lyapunov equation
10
N 100
A(T - [)J(p) + E - D]A AJ(P)A
Thecretical Density Empircall Density
0.30
Empiitical Densily
Theoretical Density
0.25
subject to Emm O, where write p- - (x), A C. and T as introduced in Proposition 4.1.
0.20
0.15
0.35
This result allows us to compare the exact asymptotic vari- ances of TD and PKL-CTD, however the forms of their variance are not as interpretable as the Monte Carlo setting. In particular, the limiting variance does not have a closed form, and must be solved for as the solution of a Lyapunov equation. There is no equivalent result for KL-CTD as we do not have convergence guarantee, however, can de- rive the asymptotic variance conditioned on convergence (see Proposition A.2).
0.30
0.05
0.00
0.00
-2
-4
-2
Figure 4. Comparison of theoretical and empirical normalized er- ror distributions at different number of updates (N) for categorical Monte Carlo across 10,000 independent seeds.
7 = 0.5 and all reward distributions supported on [0, 1]. which ensures that all return distributions are supported on [0, 2]. and then sweep over possibilities for the final support location 2m, for fixed number of support locations. We display the results from this procedure in Figure 5. The empirical results match our theoretical understanding: when Zm < 2, the approximate distribution cannot represent the full range of outcomes, resulting in biased mean estimates. Similarly, when 2m > 2, more variance is present in the Cramér projections, leading to value error.
7. Empirical evaluation
Having derived a number of theoretical insights into the comparisons between the considered algorithms, in this sec- tion we perform empirical investigation to validate these findings in practice, and to better understand the relative strengths of TD, KL-CTD, and PKL-CTD in small-scale, controlled experiments. Full details for replication are pro- vided Appendix G.
7.1. Empirical demonstration of theoretical results
Scaling with number of atoms. We next test our hypothesis
In this section, we empirically validate the theory developed in Section 6, as well as build intuition around the findings and their practical implications.
2m
10"
Asymptotic Variance. We begin by considering the asymp- totic variance results of Proposition 6.1. We empirically validate the convergence in distribution of the normalized value errors /.0/2 (V/CCC
MSE
V (x)) to its limit, and plot empirical histograms of these errors across 10,000 indepen- dent experiments in Figure 4. We find that after roughly 1 iterations the theoretical density becomes good fit for the empirical density.
16
Size of the support. Proposition 6.3 suggests that to mini- mize the increased variance due to the categorical projection, the support of the locations should be as small as possible while still containing all possible returns. To empirically study this effect, we consider an MDP with discount factor
Zm
Figure 5. Value MSE against final support location for environ- ment where all return distributions are supported on [0, 2]. Each instance is run across 1,000 independent seeds.
7

---

Categorical Distributional RL with KL Divergence
Determinsttc
Garnet Dateministic
Dirchlel Determinétic
aloms
16 atoms
10
100 atoms
200 atoms
KLOrD
12-
Cycle Gaussian
Gamell Geussian
Dirchlel Geussian
10
10
Leaming Rate
Figure 6. Value MSE against learning rate for CTD instances across varying number of locations. On the left no per-instance learning rate scaling applied, on the right learning rate is scaled by m. Each instance is run across 200 independent seeds.
Cycle T2
Gamet T2
Dirchlet T2
in Scction 6.1 that the learning rate should be scaled by vm in order to transfer across different numbers of atoms. To CTD on fixed MDP number of times
10
10"
Learing Rata
do this, with different numbers of atoms, and for each number of atoms sweep over 40 different learning rates and plot the MSE as function of learning rate. We repeat this procedure once more, but for each number of atom locations m we scale the learning rates used by / m. We present the result of these experiments in Figure 6. As was predicted by the theory, we can see that without applying learning rate scaling, larger number of support locations results in larger optimal unscaled learning rate, however after scaling by /m, the optimal learning rate is similar across a wide range of values of m.
Figure 7. Comparison of tabular TD, KL-CTD, and PKL-CTD on environments with variety of transition structures and rewards.
main (top-left), TD learning with a large learning rate is clearly the preferable approach. other domains, however, there benefits using categorical approaches. In par- ticular, in environments with high levels of reward noise, KL-CTD is preferable, and PKL-CTD performs well in all non- non-deterministie environments, indicating the benefits of an algorithm motivated by convergence considerations. We perform additional ablations over variations the algo- rithms in Appendix F.
Another phenomenon demonstrated Figure is the mono- tonic decrease in MSE obtained as m increases. This matches the behaviour predicted by the theory in Propo- sition 6.3, increasing the number of atoms for fixed support size minimizes the variance present in the Cramér projection.
7.2. Tabular experimental suite
8. Conclusion
In this section, we follow Rowland et al. (2023) and compare TD, KL-CTD, and PKL-CTD across suite of MDPs with varying stochasticity the transitions and rewards. We con- sider environments with deterministic transition structure ("Cycle"), sparse stochastic transitions ("Garnet"; Archibald al., 1995), and dense stochastic transitions ("Dirichlet"), and determinissicc/Gassiaan/ distsiiitattddddiiiuteeeed rewards. For each MDP we perform learning rate sweep for all methods over fixed budget of 1,000 asynchronous updates with the environment, and calculate the MSE of the value estimates produced by the methods.
In this paper, we have studied the fundamental proper- ties of categorical temporal -ifforal-iifferecee learning with KL- divergence. This has led us to propose a novel variant of categorical distributional RL, PKL-CTD, which makes use of preconditioning. We have proven convergence of PKL- CTD, and moreover analysed asymptotic variance of vari- ety of Monte Carlo and temporal-din differeeeeef learning algo- rithms. These analyses have led to several practical insights, including the relationship between optimal learning rates and number of categories used in approximate distributions. Natural directions for further research include fimite-sample analyses of KL--based categorical distributional algorithms.
We present the results in Figure In the deterministic do-
8

---

Categorical Distributional RL with KL Divergence
Impact Statement
Boeck, M. and Heitzinger, C. Speedy categorical distribu- tional reinforcement learning and complexity analysis.
This paper presents work whose goal is to advance funda- mental reinforeement learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
SIAM Journal Mathematics of Data Science, 2022.
Borkar, V. Stochastic Approximation: A Dynamical Systems
Viewpoint. Cambridge University Press, 2008
Borkar, V. and Meyn, S. The ODE method for convergence of stochastic approximation and reinforcement learning. SIAM Journal on Control and Optimization, 2000.
Acknowledgements
The authors would like to thank the members of the Adage Lab, Jesse Farebrother, Harley Wiltzer, Kevin Li, and the anonymous reviewers for their helpful feedback and im- provements to the paper. AMF acknowledges the support of NSERC through the Discovery Grant program [2021- 03701]. MAE was partially supported by the NSERC Grant [2019-06167], the CIFAR AI Chairs program, and the CI- FAR Catalyst grant. Resources used in preparing this re- search were provided, in part, by the Province Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.
Borkar, S. Asynchronous stochastic approximations. SIAM Journal on Control and Optimization, 1998.
Bottou, L., Curtis, F. E.. and Nocedal, J. Optimization methods for large-scale machine learning. SIAM review, 2018.
Chung, K.-J and Sobel, M. J. Discounted MDP's: Dis- tribution functions and exponential utility maximization. SIAM Journal on Control and Optimization, 1987.
Cramér, H. On the composition of elementary errors: First paper: Mathematical deductions. Scandinavian Actuarial Journal, 1928.
References
Achab, M. Ranking and risk-aware reinforcement learning. PhD) thesis, Institut Polytechnique de Paris, 2020.
Dabney, W., Rowland, M., Bellemare, M., and Munos, R. Distributional reinforcement learning with quantile re- gression. In Proceedings of the AAAI conference on artificial intelligence, 2018.
Achab, M., Alami, R., Djilali, Y. A. D., Fedyanin, K., and Moulines, E. One-step distributional reinforcement learn- ing. Transactions on Machine Learning Research, 2023.
Farebrother, J., Orbay, J., Vuong, Q., Taïga, A. A., Chebotar, Y., Xiao, T., Irpan, A., Levine, S., Castro, P. S., Faust, A., Kumar, A., and Agarwal, R. Stop regressing: Training value functions via classification for scalable deep rl. In Proceedings of the International Conference on Machine Learning, 2024.
Archibald, T., McKinnon, K., and Thomas, L. On the generation of Markov decision processes. Journal of the Operational Research Society, 1995.
Ayoub, A., Wang, K., Liu, V., Robertson, S., McInerney, J., Liang, D., Kallus, N., and Szepesvári, C. Switching the loss reduces the cost in batch reinforcement learning- In Advances in Neural Information Processing Systems, 2024.
Ferguson, T. S. A Course in Large Sample Theory. Chapman & Hall, 1996.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap. T. Mastering diverse domains through world models. arXiv, 2023.
Azar, M. G., Munos, R., Ghavamzadeh, M., and Kappen, H. Speedy Q-learning. In Advances Neural Information Processing Systems, 2011.
Hessel, M., Modayil, J., Van Hasselt, H., Schaul, T., Os- trovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M., and Silver, D. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, 2018
Bellemare, M. G., Dabney, W., and Munos, R. A distribu- tional perspective on reinforcement learning. In Interna- lional Conference on Machine Learning, 2017.
Hoffman, M. W., Shahriari, B., Aslanides, J., Barth-Maron
G., Momchey, N., Sinopalnikov, D., Stanczyk, P., Ramos, S., Raichuk, A., Vincent, D., et al. Acme: A research framework for distributed reinforcement learning. arXiv, 2020.
Bellemare, M. G., Dabney, W., and Rowland, M. Dis- tributional Reinforcement Learning. MIT Press, 2023.
http:/ http:///wwww../ www.riiutioonaleeee..
Bodnar, C., Li, A., Hausman, K., Pastor, P., and Kalakr- ishnan, M. Quantile QT-Opt for awske vision-based robotic grasping. In Robotics: Science and Systems, 2020.
Imani, E. and White, M. Improving regression performance with distributional losses. Proceedings of the Interna- tional Conference on Machine Learning, 2018

---

Categorical Distributional RL with KL Divergence
Imani, E., Luedemann, K., Scholnick-Hughes, S., Elelimy, E., and White, M. Investigating the histogram loss in regression. arXiv, 2024.
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, I.., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., Lillicrap. T., and Silver, D. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-609 2020.
Lyle, C., Bellemare, M. G., and Castro, P. S. A compara- tive analysis of expected and distributional reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2019.
Stewart, L., Bach, F., Berthet, Q. and Vert, J.-P. Regression as classification: Influence of task formulation on neural network features. In Proceedings of the International Conference on Artificial Intelligence and Statistics, 2023.
Lyle, C., Zheng, Z., Khetarpal, K., van Hasselt, H., Pascanu, R., Martens, J., and Dabney, W. Disentangling the causes of plasticity loss in neural networks. arXiv, 2024.
Sutton, R. S. Reinforcement learning: An introduction. MIT Press, 2018.
Meyn, S. Control Systems and Reinforcement Learning. Cambridge University Press, 2022.
Székely, G. E-statistics: The energy of statistical sam- ples. Bowling Green State University, Department of Mathematics and Statistics Technical Report, 3(05):1-18, 2003.
Morimura, T., Sugiyama, M., Kashima, H., Hachiya, H., and Tanaka, T. Nonparametric return distribution approx- imation for reinforcement learning. In Proceedings of the International Conference Machine Learning, 2010a.
Székely, J. and Rizzo, M. L. Energy statistics: A class of statistics based on distances. Journal of Statistical Planning and Inference, 1249)))249-12777 2013.
Morimura, T., Sugiyama, M., Kashima, H., Hachiya, H., and Tanaka, Parametric return density estimation for reinforcement learning. In Proceedings of the Conference Uncertainty in Artificial Intelligence, 2010b.
Thakoor, S., Rowland, M., Borsa, D., Dabney, W., Munos, R., and Barreto, A. Generalised policy improvement with geometric policy composition. In International Confer- ence on Machine Learning, 2022.
Peng, Y., Zhang, L., and Zhang, Z. Statistical efficiency of distributional temporal difference learning. Advances in Neural Information Processing Systems, 2024.
Virtanen, P., Gommers, R., Oliphant, T. E., Haberland, M., Reddy, T., Cournapeau, D., Burovski, E., Peterson, P., Weckesser, W., Bright, J., van der Walt, S. J., Brett, M., Wilson, J., Millman, K. J., Mayorov, N., Nelson, A. R. J., Jones, E., Kern, R., Larson, E., Carey, C. J., Polat, i., Feng, Y., Moore, E. W., VanderPlas, J., Laxalde, D., Perktold, J., Cimrman, R., Henriksen, I., Quintero, E. A., Harris, C. Archibald, A. M., Ribeiro, A. H., Pedregosa, F. van Mulbregt, P., and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272 2020. doi: 10. T0.1038/41155 090-5599---0-66666-
Peng, Y., Jin, K., Zhang, L.,, and Zhang, Z. Finite sample analysis of distributional td learning with linear function approximation. arXiv preprini arXiv:2502. 14172, 2025.
Puterman, M. L. Markov decision processes: Discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Robbins, H. and Siegmund, D. A convergence theorem for non negative almost supermartingales and some applica- tions. In Optimizing methods in statistics, pp. 233-257. Elsevier, 1971.
Wang. K., Zhou, K., Wu, R., Kallus, N., and Sun, W. The benefits of being distributional: Small-loss bounds for re- inforcement learning. In Advances in Neural Information Processing Systems, 2023.
Rowland, M., Bellemare, M., Dabney, W., Munos, R., and Teh, Y. W. An analysis of categorical distributional re- inforcement learning. In International Conference on Artificial Intelligence and Statistics, 2018.
Wang, K., Kallus, N., and Sun, W. The central role of the loss function reinforcement learning. arXiv, 2024a.
Rowland, M.. Tang, Y., Lyle, C., Munos, R., Bellemare, M. G., and Dabney, W. The statistical benefits of quantile -ifferee-d learning for value estimation. Proceedings of the International Conference on Machine Learning, 2023.
Wang. K., Oertell, O., Agarwal, A., Kallus, N., and Sun, W. More benefits of being distributional: Second-order bounds for reinforcement learning. In Forty-first Interna- tional Conference on Machine Learning, 2024b.
Rowland, M., Wenliang, L. K., Munos, R., Lyle, C., Tang, Y., and Dabney, W. Near-min meam-minimax..o optimal distributional reinforcement learning with a generative model. arXiv preprint arXiv:2402.0759, 2024.
Wu, R., Uchara, M., and Sun, W. Distributional offline policy evaluation with predictive error guarantees. In Proceedings of the International Conference on Machine Learning, 2023.
10

---

Categorical Distributional RL with KL Divergence
Wu, W., Li, G., Wei, Y., and Rinaldo, A. Statistical infer- for temporal difference learning with linear function approximation. arXiv preprint arXiv:2410.16106, 2024.
Wurman, P. R., Barrett, S., Kawamoto, K., MacGlashan, J., Subramanian, K., Walsh, T. Capobianco, R., De- vlie, A., Eckert, F., Fuchs, F., Gilpin, I.., Khandelwal, P., Kompella, V., Lin, H., MacAlpine, P., Oller, D., Seno, T., Sherstan, C., Thomure, M. D., Aghabozorgi, H., Bar- rett, L., Douglas, R., Whitehead, D., Dürr, P., Stone, P., Spranger, M., and Kitano, H. Outracing champion Gran Turismo drivers with deep reinforcement learning. Na- ture, 2022.
Yang. D., Zhao, L., Lin, Z., Qin, T., Bian, J., and Liu, T.-Y. Fully parameterized quantile function for distributional reinforcement learning. Advances neural information processing systems, 32, 2019.

---

Categorical Distributional RL with KL Divergence
APPENDICES
For convenience, we collect together the contents of the appendix:
Appendix A contains several additional results referenced in the main text.
Appendix contains the proofs of results stated in Section 4.
Appendix C contains the proofs of results stated in Section 5
Appendix contains the proofs of results stated in Section 6.
Appendix provides a summary of several areas of related work.
Appendix F provides additional experimental results referenced in the main text.
Appendix G contains full details for all experiments presented in the main paper.
A. Additional Results
A.1. Counterexample to KL as Lyapunov function for KL-CTD
We construct numerical example that shows that even in single-state MDP, the KI. divergence KI. (x) p (x)) need not decrease under KL-CTD dynamics. The example is constructed follows.
We take single-state, single-action MDP, such that taking the action in the state leads to a ceff-transition. We take the immediate reward to be 0.02644103, and discount factor to be 0.99. For the categorical support, we take 10 equally spaced atoms between and 100. We consider initial values for the logit vector given by
(1 (1.88191424, 0.02041108 - -1.22448044444 0. 0.44203928, 0.71425795. 0.46704711, 0.09271942, 0.1.2,0,...11111 0.11709652, -..32497122, -1.5971562))
After taking softmax, this yields initial categorical probabilities of approximately
(0. 40581378, 0.06055603, 0.01816506, 0.09616057, 0.12624672,
0. 0.09859566, 0.06780931, 0.06948262, 0.0446569, 0.01251333);
see Figure 8 for an illustration of these initial categorical probabilities Po. the target distribution Tr under the categorical distributional Bellman operator, and the fixed-point categorical probabilities 7*
0.4
0.3
0.3
0.6
0.2
0.4
0.2
0.0
0.0
40
60
80
100
40
60
100
20
40
80
100
Figure Left: The initial categorical probabilities po. Center: The initial target probabilitics, T" po Right: The fixed-point categorical probabilities, p
We simulate the cumulative effect of expected KL-CTD updates with small learning rate by numerically solving the flow
0. (7 m)
using the default scipy- integrate. wwwve.ilviiiveee method (Virtanen et al., 2020). We plot KL(P" ( against the ODE time t in Figure 9, and note that initially, the KL divergence increases, meaning it is not a Lyapunov function for the KL-CTT dynamical system.
12

---

Categorical Distributional RL with KL Divergence
0.85
0.8
0.84
0.6
KL Ipp)
KLip "Ipp)
0.83
0.82
0.81
0.2
0.80
50
100
150
200
1000
2000
3000
4000
ODE time
ODE time
Figure 9. Numerical illustration of the non-monotonicity of the KL divergence for KL-CTD dynamics.
A.2. Lipschitz coefficient of the KL gradient
Proposition A.1. The gradient V KL[g ||
softmax((4) is Lipschitz with constant at most
Proof. We note that V KL || softmax(p)) = softmax(s)) - and we also note the Jacobian of this is given by J(softmax(()), where we use (p) = diag(p) - ppp We next recall the fact that function being Lipschitz with constant equivalent to all eigenvalues of its Jacobian being less than every point in its domain. We also recall the fact that the Frobenius norm of matrix upper-boundis its largest eigenvalue. We now consider
sup ||]("))|=== Fm
sup Rm
Σ (1 p2)2 + Σ (1F))(3))
m
m²
as the supremum is attained when J(softmax(4)) is a uniform distribution. Combining with the previous results, this implies that all eigenvalues of the Jacobian upper bounded by y/-I and complete.
Proposition A.2. Suppose the iterates (TCD)) (OCD))20 322 generated by following Equation (6) with step size ag agk -8 for
CTD 3 (1/2, 1), and moreover suppose that this sequence converges. Let us write V.TTD = T p
Then we have that
(x) - (z)) 3 NO TT (()"(())) (c)))
where EcTD(æ) is the unique solution of the Lyapumov equation
(T
>J (4)) - J(p")
subject to 1T1 0.
Proof. We first show that the logits must converge to a unique point. From Proposition 4.3 we know that we must have convergence
o such that (x) pd (z) for each æ, which gives a 1-dimensional subspace of solutions for each x. However we also know that each state the CTD update orthogonal 1 (x) e. (z)) 0, which uniquely identifies o.. Since under assumption we have convergent sequence to unique fixed point we satisfy the assumptions of Proposition D.1, which gives us that
so that must have
3/2 (0 (x) - cx(=)) 4 (TD,))D(E)),
where ETD(F)) is solution of the I Lyapunov equation
(T - ()) SJ -I -
This equation admits a 1-dimensional subspace of solutions, however we know that there is no covariance in the direction parallel to 1, so we can uniquely identify the solution through the constraint TT = 0. The remainder of the result now follows directly from applying Proposition D.2 to the map @ + T p and we are complete.
13

---

Categorical Distributional RL with KL Divergence
A.3. Extension asynchronous PKL--TTD
Our convergence result in Section 5 focused on the case of asynchronous updates; here we discuss a high-level how the key details map onto settings with asynchronous updates. In particular, from sequence of transitions (Xx+ RR x20 compute sequence of logits (Øx k220 given by
(x+1(x)= =x( aRe C
Σ th (X)) + y2s) - =xx()))
for x = XRs (z) 0x (x) for * Xx. and ag is state-dependent stepsize.
In this setting, under certain assumptions on the step sizes and distribution of updated states (Xx)a20. the ODE becomes
0.0((z) = c(x) CTC -Ip
(14)
where c(x) is a constant reflecting the relative update frequency of different states. It can verified that these dynamics converge, using a Lyapunov function argument similar to that of Proposition 5.4.
Proposition A.3. The ODE in Equation (14) converges, and a Lyapunov function for its convergence is
d" (x)
I(Ø) = Σ
KL(
(1:) (x)])p(x)))
c(z)
NEX
Proof. Writing the right hand side of the ODE ((),, we can write out (VoI()),
d
(VoL()).(())
c(x) (p (x) C(T"')))))))(() (x)))
Σ d= (x)
which becomes the same expression as Appendix C.2, which is guaranteed to be negative.
complete asynchronous convergence proof from here relies the details on the sequence of transitions (Xa, Rx. X1) (e.g. if (X forms ergodic Markov chain), and details the step size sequence. Suitable results which can be used as basis to build upon can be found in Borkar (1998) and Borkar & Meyn (2000).
B. Proofs of results from Section 4
B.1. Proof Proposition 4.1
Proposition 4.1. The expected KL-CTD update can be expressed via malrix-vector multiplication as
& с (TT T)p²
where we interpret & and pd as |X| x m-dimensional vectors, and where T is the categorical distributional Bellman operator (Rowland et al., 2018; Bellemare et al., 2023; Rowland et al., 2024), given by
T (x, i; y.j) = P" P" (z) ( (B
Proof. To begin, recall that the synchronous sample-base based KL-CTD update takes the form
((x)= +- <xx) - Σ
+ + || (+()))
where for each state X, (x, R X*) an independent transition sampled under и. Next, note that for vector yp E Rm and distribution q E Am, have
WLL(g softmax(p)) softmax(p)) q.
14

---

Categorical Distributional RL with KL Divergence
As result, the expression for the update above can be rewritten
<(x) ((x) + E P. (X*) (N(R² + - p0()))
Now, taking an expectation over the random transition (z, Rt, x*) yields
E
Σ p (x*)
- (0))
M
p" () 722)1 - (x)
1= MEN
Pd ( - (x))
1=1 yEX
= =( - (x),
as required.
B.2. Proof of Proposition 4.3
Proposition 4.3. The only stationary point of expected update in Equation (7) is the fixed poins 5e of T
Proof. At a stationary point, the expected update is O. From Equation (7), the expected update (up multiplication by the learning rate takes the form (T - T)p and thus the expected update being 0 implies T po p From the contraction theory recalled in Proposition 4.2, T has fixed point, p", and is also contractive map on A does not have any other fixed points in m Thus, at stationary point, we must have p fp" as required.
C. Proofs of results from Section 5
C.1. Proof of Proposition 5.3
Proposition 5.3. The mapping T A > A m is a contraction with respect to the norm || ||=: for any P. q E X we
have
1 - T I V7llp All.
where || ||x is the norm defined by Ilplxx = EvEx d (=<)\((p())ll
Proof. Let 1 E P(R)). We follow the structure of the proof of contractivity of T in the norm || described in Proposition 4.2, and established by Rowland et al. (2018). Thus, begin by observing that by identifying A.. with P((=i,..,, }) via the map m + >m prO2. the operator 7 can be interpreted the composition of the distributional Bellman operator T (( 1) + (R) and the categorical projection Ile Rowland et al. (2018); see Bellemare et al. (2017); Morimura et al. (2010a); Chung & Sobel (1987) for earlier definitions of the distributional Bellman operator alternative forms. We define the r-averaged version of the Cramér distance (see Proposition 4.2 as
C(an') Σ ((())(l)))(()())) xEX
15

---

Categorical Distributional RL with KL Divergence
for all n/4/: X- P (RR).
Then, arguing in Rowland al. (2018), have
TTTT,
Σ
d' (x)
(T"
NEX
[ Σ (br)))((X')) X' X1
/ Σ (brr)mu(X")) Ps (dr, X' 12)
Σ
d' (()) l
X
SEX
X'EX
< Σ (()
E
P" X' x)) ((br.)) (b),)))))))))))(()))))))) (brn))u(X())
12)
REX
JB X'EX
=7 E (x)
E
Pa (dr,, X' X' & ((xx')(((X)))
ZEX
B X
E d'(x') l2 (n(x')
= Y
X'ex
7 12 (2)
Here, (a) follows from the invariance of d", used in the argument of Wu et al. (2023) in the case of 1- I-Wasserstein distance. Further, Ieç is non-expansion in & (Rowland et al., 2018), so together, we have that II is a cn-cottraction (( }m with respect to f distance. Further, the map between Am and ({ ,m) }m) is an isometry when Am is equipped with the metric that sets the distance between p. E Am to be ||C(p q)ll2, and (1 }m) is equipped with lz. Hence, we recover that T a V--ootractionn on A with respect to || l| as required.
C.2. Proof of Proposition 5.4
Proposition 5.4. The continuous-time dynamics of Equation (10) converges, and a Lyapunov function for its convergence
I(0) = Σ KL (
(1))
ZEN
where d= the stationary distribution of the policy w.
Proof. We can note that the gradient V V (()) is given by
VoL(Ø) = -z) -(76)
BEX
Writing the right hand side of Equation (10) as f(0). we can now write out (V ((L())
V VoL(O), = (E)
C(T'' -p*())))
Σ
[(14) -
T
Σ
-
(+)
X
16

---

Categorical Distributional RL with KL Divergence
We note that
ZEX
mEX
Σ
(FF (
(- (1))) [[)
ZEN
114
Σ
d')x) (
d" (<) - p p
LEE
WEX
Edd(s) (x) [C(p"(t) - -(0)))1 1
V7
ZEN
xE
V ZEN
(
where both inequalities marked (a) follow from the Cauchy-Schweer inequality, and (b) follows from the contractivity established Proposition 5.3. This shows that the right hand side of (+) strictly negative when of + pnd hence we are complete.
C.3. Proof of Theorem 5.5
Theorem 5.5. Suppose that (ox x20 is a sequence of logits generated according to synchronous PKL-CTD updates. That is, for each O, we have independent transitions (Ir, RR: Xx) such that
xx+1(x) = (x)xx) xxxx
CTC
Σ p (XE)A(RE + +22) p (1))
Further suppose that the stepsizes (0x))20 satisfy the Robbins-Muune Munro conditions En-n Og and Eee of < 0o. Then we have that db converges in the sense that you (x) +0*(x) for every x E X almost surely.
Proof. The high-level structure of the proof follows the classical Robbins-Siegmund theorem (Robbins & Siegmund, 1971), using the fact that as established in Proposition 5.4, the function L in Equation (11) is Lyapunov function for the associated continuous-time dynamical system, and combining this with the supermartingale convergence theorem. We begin with the assumption that d stationary distribution of n, has full support over the state space X. We will then treat the more general case where this assumption does not hold by following the argument made in Thakoor et al. (2022), by making inductive chaining argument across the communicating classes of the Markov chain induced on X by n.
To begin, we perform a second-order Taylor expansion of the Lyapunov function L((x+1) around Φx. using Eg to denote conditional expectation given the random variables defining the sequence (Ø1)osigi, and pie to denote pdx:
E, [L(Φx .)] E& [L C C(T - T)ps + aHEN)]
L(dx)) an (VL())), C(T 1)p2)
+ a CCT" - 1)pk + Ek. (L( C(T" - 1)pk + apEn)]
S <((b)) (1 / -plll
o Ex [CT C(T T)pk ck, HT((x) C(T T)pp + amex))
where ((v)) is the Hessian of L at db. ex(x) CT the zero-mean noise at step k. and the inequality follows from the argument in the proof of Proposition 5.4. The Hessian
1)x (z) - -
+ + p. (xt)
17

---

Categorical Distributional RL with KL Divergence
of the Lyapunov function is uniformly bounded, and too the quantities C' C(T Dyp Eg. there exists a global constant B 0 such that
+ a... Ex [[()++1]] S (n -i)-Ø>
This is almost a supermartingale inequality, save for the final term on the right-hand side. The idea, however, that thanks the summability condition Ex20 aj < co, the cumulative effect of these terms does interfere with the convergence guarantees associated with non-negative supermartingales; the following argument follows the approach of Robbins & Siegmund (1971).
Writing La = I(ox), note that if we define
k-1
La Lx -
Σ
aj B
Σ
1=0
1=0
then we have
Eg [R[[+++]= Ex [1k+1]
Σ
2 +
Σ
on(1
1T)l - -P 1
1=0
1=0
k
< Lx -ax(1 - vT
pall + agpb - Σ
ai + Σ
a1 v>) - pill
1=0
1=0
aips
ar(1 v -
L
Hence, (I=)a20 is a supermartingale, and uniformly bounded below by -pk Ex=o az and by the martingale convergence theorem, we deduce that Lx converges almost surely. therefore follows that
Σ ca - 1=0
converges almost surely, and hence so does LA- If La does not converge to 0 almost surely, then there exists >> such that lim inf f Lg > e with positive probability, On this event, we would therefore have that
diverges, since ΣeS ay 0o, contradiction. Hence, La 0 almost surely, and this implies pk - pr almost surely, as required, since we assume d" has full support.
To make the general argument, where d does not have full support, we can proceed in exactly the same manner as Thakoor et al. (2022). First, by adjoining a terminal state necessary (to deal with environments in which episodes terminate in finite time), there exists an invariant probability distribution d= over the state space X' under the transition dynamics induced by x. The argument above applies verbatim to obtain || p**1 | + (though dt does not have full support, this does not show Pk ).).
Now, we apply inductive argument to the communicating classes of the Markov chain on X' induced by n. The set of communicating classes of finite state Markov chain form a directed acyclic graph (with a directed edge drawn from one communicating class K2 to another K₂ if there exist x E K1,y y E K₂ such that P (y]az) > 0. The leaves of this directed acyclic graph are precisely the communicating states which have an invariant probability distribution. For any such invariant class K, there exists an invariant distribution d supported on K, and the argument above applies to obtain pk (x) - 57 (s) for all æ E K, and moreover, that
Σ cn(1 - prll=
18

---

Categorical Distributional RL with KL Divergence
converges.
We now induct communicating classes, based on the maximal length of path in the directed acyclic graph from the theory there exists a probability
communicating class to a leaf. For any such communicating class K, by distribution K supported on K, with the property that
RP = AE + (1 A)w'
with À E [0, 1), and probability distribution supported solely on communicating classes K' for which, by inductive hypothesis, we have px(x)) + (x) for all E K". Now, following exactly the steps of the proof of Proposition 5.3, we obtain
17Tp - TT 912 I Ayllp فll + + (1-x)llp g
where semi-norm || Ilx is defined by 1|p|~ - E <(z)]lp(z)) and similarly for || llk..
Using the inequality above with p = 1" and Pk. and from the expansion
11 Tst pall = lp - pxll + llpk - T pxll2 + 2( - -kk. C CTp(p T-' px))a
where U)x = ExEx
v())) (u, E x denotes the w-weigtted inner-product, we obtain
- + lips Pxlli + 2(p" pk. C C(p px)) I Ayllø" p + (1 \illø" -00--palle P
which, with some rearranging, yields
TT-
Pk: C C(pp T px)) S ((Ar 1)115" pall2 + (1 A) -p& 11 1m 11:)
Now, defining
I²(0)= Σ M(4)KL((777()11))0)))
SE.
and writing LE (dx),, we can compute that
Ex 1 S LE + (VL) C Pk -pk) a²B
= - + xx - Pk: C C'C(p: PTpp))): ) + +
\y 1pt - Pk 1l + aa 1 X
S L + 2 B + a
11p
-P P
By the induction hypothesis, Σ 1= 04 1pp prll? converges, and may apply the supermartingale convergence theorem
as above to conclude that Ei a/ / pT' - pr ? converges, and that -|pT - ps llx - O, completing the inductive step and the proof.
D. Proofs of results from Section 6
D.1. Proof of Proposition 6.1
Prior to the proof of Proposition 6.1, we use the following supplementary results.
Proposition D.1. (Meyn, 2022, Theorem 8.1, Proposition 8. 8.10) Let f IRd - Rd, and for an initialization 0g E Rd. consider the sequence of iterates (6m)m20 defined by the iterative rule
0a+1 02 + an+11f(00))) - En),
with (en)n20 a mean-zero martingale noise sequence, such that the distribution of depends only on On and oz, = a/n/. If (i) 0,0,- ((.) has a globally asymptotically stable equilibrium 8*, (ii) 0₂ 0* almost surely, (iii) f Lipschitz continuous, (iv) V f(") is Hurwitz, then the scaled error Om 11/2 (0, - 8*)) converges in distribution to N(0, E), where Σ is a positive-defintte solution of the Lyapunov equation
AΣ ΣAT = -EAs
where EA the covariance of the noise en at 0 and A (f(()).
19

---

Categorical Distributional RL with KL Divergence
Proposition D.2. Let g: Rd - R be a mapping such that the Jacobian Jg is continuous in a neighbourhood of p E Rd, If Xn is sequence of d-dimensional random vectors such that for some B3 [0, 11, nD (Xn - -1² S, E,,, then
n (g(X,)) - 9(()) 4 (((0JJEJJJ)
Proof. This proof closely that of Theorem 7 of Ferguson (1996), but adapted to the scaling na instead of vn. We note that we have Xn, - J.. Then if there exists 8 > 0 such that g is continuous on the set {or E R || pll2 < 8), for such that | - pll2 <: 8 we have
= g(p)+ + S// Jolu+ v(x de (2- 1)).
So for such that |xxn plll < 8,
n (g(x)) - g(1)) = =n S( Jzlil + u(Xx - 11) du (x - 1t).
Since X, 4, ,, P(IXX, 11l 6) + and Jg(H -
de - so n (o(xx) 4 Jg(N)) where NOO, (لاير and we complete.
Z- (0, E). From this we note that (J
Lemma D.3. Suppose the iterates ( 220 and (6FCMC))>)) )x20 were produced using step size ag =akk for 3 E (1/2, 1)
)x20
and P. (x) > for each x, E X' x [m]. Then we have that there exists unique fixed points &cMC and @. PPMC such that have
k0 (o
CMC (x) -0. (x)) S (0,EEMM))) EcMc)
(x)
and
PCMC (x)) $ (0, ErcMc), (Ø (x) o0
PCMC
where EcMc= = [I - m 11 and ΣPCMC c1
Proof. We begin by deriving the fixed points @CC and dPCMC, To begin, we note that from the assumption on full support of j each state, the function ((x) log(f" (x:) satisfies softmax(0(x))
(x) for each x. Furthermore, (oftmax(0(x)) + cl) (x) for any c, and this construction contains all logits satisfying this condition. We write s -= = X + R Be E IR. ((x) d(x) + cl) for this set of optimal logits, and note that each 6 E $ is fixed point for both CMC and PCMC. We now show that CMC and PCMC always converge to a unique point in this set, regardless of the given sequence of updates.
The CMC update at each step and state is orthogonal 1, as consists of a difference of probability vectors. This gives us that beginning from of @CMC, for any k - 0 we have 1 T (MCC(() of CMC (x)) 0, for any x E X. From this we can 0 for all x E X' and remark that this is entirely dictated by @CMc. The PCMC updates also maintain an invariant in their updates: in particular the updates do not change the mth element of the logit (this can be seen as (C' Ce)m = 0 for any e E 4). Thus o." PCCC is the unique element of & such that
conclude that &CMC; is the unique element of t such that 1 (oCMC()) (a) % (CMc(x)
(x)) = d
CMC (x)m for all x E X.
We may now apply Proposition D.I, since CMC converges as stochastic gradient descent on a convex objective with bounded noise, and the convergence of PCMC is guaranteed as a special case of the analysis carried out for PKL-CTD (with discount 0) in Theorem 5.5. The existence of a Lyapunov function for both systems guarantees the global asymptotic stability required as well the Hurwitz property for the Jacobian at the fixed point. Using Proposition D. it now remains to find solutions to the I Lyapunov equation of each method. Before applying the lemma, we pause ensure that all of its conditions are met. The previous arguments gave the unique equilibrium point for both CMC and PCMC, and the remaining assumptions are met if the methods converge. We note that CMC converges as it corresponds to stochastic gradient descent on a convex loss function, and we also have that PCMC converges as it corresponds to PCTD with a discount factor of O, which we have convergence for (Theorem 5.5).
Now applying Proposition D.I 1 to CMC, this corresponds to solving
J
20

---

Categorical Distributional RL with KL Divergence
at each state X. As J(p (x)) is full rank (the vector 1 an eigenvector with eigenvalue 0), there is a I-dimensional space of solutions. However also know that the CMC update is orthogonal 1, the asymptotic variance is the unique solution E' of the Lyapunov equation subject 1 Σ1 = O. We now show that this matrix Σ is given by (L 11 )) To see this we can note that Col(11 span(1) = Ker(J(pp (z))). so that J(" (x)) J (a)). We also note that as [["" (ar) and 11 T are both symmetric we have that 11 T [ (x)) = O, and hence EJ( = 1 (f)")))) (I)). Combining these we have that
2
(x) EJ ( J(
as desired.
We next turn to PCMC. This now corresponds to solving
C' CT - CCCCC)C))C)))CCC
As in the previous case, this equation admits a 1-dimensional subspace of solutions, however from the assumption above we know that we must have O covariance in the bottom right element of Σ, as there is no change in this logit. We can then find that ((T C- C--11 111 satisfies this condition, and we can verify that this satisfies the Lyapunov equation.
Lemma D.4. Under the same assumptions as Lemma D.3, we have
(x) (x)) 4 N (0. 1 (1)))))))
and
(z) AN N (0. ( (0.J0))CTcC6)))
where the shorthand J(p) diag(p) pp T.
Proof. This follows from the convergence result of Lemma D.3 and Proposition D.2. In particular we are applying the function σ(z) + (x). which has Jacobian J(p((x)) (x:)). We can also note that 1 is in the kernel of (xp(x)) so that
(l - 11 (1 ) (x))
and
J(
With this, we are complete.
Proposition 6.1. Suppose the iterates k20, (oEMC) )220 and (Ø (PPMM)) k20 were produced using step size og = aok for 3 E (1/2,1 1). Further suppose that 8. (x) > 0 for each x, i E X' x [m]. Then we have that
k20,
1 (V V*(x)) N (x)) k (8 CMC dF
(x) VT())) SN (0. 2 ((F((()))
1 (8/ (x) - Vx)))) $ N (0, fu()) u(x))
where (p) diag(p) pp u(z) = C)( and o² is the variance of ntt (x).
Proof. This follows from the results of Lemma D.4 and Proposition D.2, as we are passing the converging probabilities of Lemma D.4 through the map + TT, 2, and the effect on the limiting distribution from this transformation given by Proposition D.2.
21

---

Categorical Distributional RL with KL Divergence
D.2. Proof of Proposition 6.2
Proposition 6.2. Suppose p- (pi)i= is collection of probabilities associated to the locations (zi) i=1 and let the variance of this categorical distribution be o². Then there exists 3 depending only on p such that
Bo² < <2 J(p)222>>>
Proof. Note that J(p) is positive-semideffiite matrix with all eigenvalues in [0, 1) (we have T J(p) 2 0 for any it is the variance of random variable distributed as ,, P:d., and its maximal eigenvalue must be strictly less than 1 as its trace is P. - p:) < _ Writing 0 = A - -... I Xm < for the eigenvalues of [p), we note that .i...., An are the eigenvalues of J(p) and J share the same respective eigenvectors as J (p) is symmetric. Next writing 2.....; 2m for the coordinates of in the basis of corresponding eigenvectors J(""), ), we have that
2 JT)p==== = Σ
2 A
and
m
2 J(p)²² = Σ x?x?.
As each X < 1, this results in strictly smaller quantity. We can also note that 2 J(p)z s Am and 2 [(p)22 2 X2
>3|12112 so that x J(p)s and we can set to satisfy the statement in the proposition.
J(m)
Am
D.3. Proof of Proposition 6.3
Proposition 6.3. Suppose v is a probability measure whose support contained in [a, b), and Varz 02. Let
21,.... :m be equally spaced points, Then we have
Varz~Hew (2) = = 20
(b-a)² where 0 S B(v) < 1m- is quantity capturing the amount of projection in the mup + Ngv.
Proof. Let Z ~ v. then by the assumption on the support of v we have that 21 S Z S am almost surely. Our goal is to introduce a jointly distributed random variable E distributed as the difference between Z and the Cramér projection of Z.
We define ((Z)= = max(z; 24 24 S 3}. and u( ()) = min(2 (2: : 244 2 }}. If (ZZ) = u(2) then Z is already on a support point of IICR and in this case we set E to be with probability 1. In the case that ((Z) < (Z)), we set E to take value (Z ((2)) with probability (u(Z)
(u( -(Z)). and (u( (u( - Z) with probability (Z (ZZ)((Z)/((uf) (2)-Z))))) Then we have that Z + E takes value I(2) with probability (u(Z)- /u)(((((2)) ((Z)) and u(Z) with probability [Z (z))))(((2)-1122) that is, Z + E is distributed according to 1lg..
We can note that E[E| 2| = O, which gives us that E and Z are uncorrelated. We can also consider Var(E) (E 2), which therefore reduces to computing B[F22 Z]. When Z lies on the grid, this quantity is trivially 0, and otherwise we have
(u(Z) Z) (u(Z) - ((2)) + (u(Z) 2)² (u(Z) [(Z))
(Z ((2)
E E2
2 = (Z - 1(
(Z ((((())(((Z) 2) (u(Z) 2) + (2 (11)
(u(2) - ((2))
(Z ([(())((((Z)) 2).
22

---

Categorical Distributional RL with KL Divergence
We next set U Z + E, and write
Var(U) Var(Z) Var(E)
Var(Z) + IE [Var( E 2)] + Var (E [ Z))
Var(Z) + E (aa E Z)]
(Vr(Z)+ Σ
Var(E = == P(dz)
m-1
-ar(ZZ)
-2 (z ( P(d2))
where (a) follows from the deduction above that Z and F. are uncorrelated. This quantity is minimized when v is supported on (z1,..., 2m }. as the additional variance term is O. The quantity (z (==1(=)) (==) - =) is maximized when = (=)))) meaning 2 is concentrated on the midpoint of 2; and 4,+1, that is, the case when Ei= pro(acta++1)/2- and the amount of projection is maximal. In this case, the additional variance incurred is given by
(u (2) - 1()))2
(b a)²
Σ
4
4(m 1)²
1=1
(b
a)²
4(m 1)²
D.4. Proof of Proposition 6.5
Proposition 6.5. If the step sizes (ox) are given by ag = ook- we have that
- AN (0, by ΣPCTD(I)b.)
where by = J(i and (PCTD(F) is the unique solution E of the Lyapunov equation
A(T - T)J()) 5.0
AJ(PA
subject to Emm = 0, where we write = (z). A = C. and T as introduced in Proposition 4.1.
Proof. We may apply Proposition D.1. Our earlier Theorem 5.5 guarantees the conditions of the result hold, thanks to the establishing of Lyapunov function; the existence of the Lyapunov function establishes global asymptotic stability of the equilibrium, and the expected update function clearly Lipschitz (since as noted earlier, the softmax function is Lipschitz), and the Lyapunov function also guarantees that the the Jacobian at the equilibrium point Hurwitz. We thus obtain the stated Lyapunov equation for the term MPCTD. This equation will have a 1-dimensional subspace of solutions due to the rank-deficiency of (p), however the PKL-CTD updates do not modify the final logit, we must have the bottom right entry of the covariance matrix be O. This constraint uniquely identifies the matrix in the subspace of solutions as the subspace of solutions is along the span of 11 T_ This then combined with Proposition D.2 to reflect the effect of the map + z dd completes the statement.
E. Related work
Categorical approaches tasks that are traditionally modeled as regression have been studied in variety of settings more generally. Stewart et al. (2023) study the effects on neural network representations, and Lyle et al. (2024) study the effects on network plasticity.
Beyond categorical distributional reinforcement learning (Bellemare et al., 2017), a variety of categorical approaches to reinforcement learning have been considered the literature, including two-hot regression targets, and various versions
23

---

Categorical Distributional RL with KL Divergence
of one-step distributional RL (Achab, 2020; Schrittwieser al., 2020; Hoffman al., 2020; Achab al., 2023; Hafner al., 2023). In partieular, Farebrother al. (2024) study the empirical benefits of categorical algorithms for reinforcement learning, finding combination of one-step categorical distributional reinforcement learning (Achab, 2020; Achab et al., 2023) and implicitly injected Gaussian noise to smooth the resulting loss (Imani & White, 2018; Imani et al., 2024) particularly effective, producing a novel reinforcement learning algorithm that was demonstrated to be performant in a wide variety of applications.
The chamérvvariant of CTD was first analysed by Rowland et al. (2018), who proved its convergence to the fixed point of the categorical dynamic programming algorithm. Boeck & Heitzinger (2022) then developed variant of this algorithm that incorporates ideas from speedy Q-learing (Azar et al., 20120 1); both Boeck & Heitzinger (2022) and Peng et al. (2024)
obtain finite-time sample complexity bounds on Cramér-based CTD algorithms. An analysis of a preconditioned variant of Cramér-CTD was performed by Peng et al. (2025).
Rowland et al. (2023) study the efficacy of quantile remporal-difference learning (Dabney et al., 2018; Rowland et al., 2024), a complementary approach to distributional reinforcement learning, for value-function estimation. Several recent works also analyze the complementary class of wowllloodd-axmmizzttini based methods for distributional reinforcement learning, both from the perspective of regret minimization (Wang al., 2023; 2024b; Ayoub et al., 2024) and sample-efficient policy optimization and evaluation (Ayoub al., 2024; Wu et al., 2023; Wang et al., 2024a).
F. Additional experimental results
We perform ablations over potential causes for the occasional underperformance of TD: synchronous vs asynchronous updates, boundedness of the value estimates, and boundedness of the regression targets. The experimental suite in Section 7 uses asynchronous updates for TD, KL-CTD, and PKL-CTD, which may cause an increased difference in the number of updates across states. We present the experimental suite using synchronous updates in Figure 10. The second potential cause is boundedness of the value estimate: due to the nature of the categorical parametrization, the categorical value estimate at each state is bounded in [21, ]. To see if this has material effect on the relative performance of the algorithms, we consider a variant of TD learning where the value estimate at each state is clipped to [z], =m] after each update. We present the experimental suite under these updates in Figure 11. The third potential cause is that the categorical target distributions are always bounded in [Zi =m.. which may be contributing to a regularization effect, especially with heavy-tailed rewards. To test this we consider a variant of TD learning where the TD target is clipped to [21, 2m] before being used in the update step, and present the experimental suite using these updates in Figure 12.
Across all of the above considered ablations, the relative performance of TD, KL-CTD, and PKL-CTD are fairly robust, suggesting that in the settings where KL-CTD or PKL-CTD outperform TD, this is more likely to be result of the underlying algorithmic differences, rather than an effect of the relative number of updates per state or clipping effects the categorical parametrization.
24

---

Categorical Distributional RL with KL Divergence
Cycle Delerminittic
Garnel Deleeminisiic
Dirichiel Deterministi
TD
KLLOTD PRUIID
10"
Cycle Gaussian
Gamell Gaussian
Difichil Gaussian
16'
Cycle T22
Gamel T2
Dirichee T2
16
10"
12"
10/
10"
12'
"0"
16'
Leaming Rate
Figure 10. Experimental suite where each algorithm is run synchronously.
Cycle Delerministic
Gamel Deterministic
Dirihhell Deterministic
TO
KL-CTD PXQ.-CTD
10"
Cycla Gaussian
Gaussian
Dirihhill
1200
Cycle T2
167
167
Leaming Rale
Figure 11. Experimental suite where TD estimates clipped to [21, 2m) after each update.
25

---

Categorical Distributional RL with KL Divergence
Cycle Deterministi
Gemell DDeterinistic
Dirihiell Deterministi
16
TD
KL-CTD PAU-CTD
10
Cycle Gaussian
Gamel Gaussian
Difihhel Gaussian
1//
Cycle 72
Gamel T2
Dirichett T2
Loaming Rate
Figure 12, Experimental suite where TD targets clipped [21, 2m) prior each update.
26

---

Categorical Distributional RL with KL Divergence
G. Experimental details
In this section, we collect full details for the experiments presented in the main paper. For all generated figures, the error bars represent =2 times the sample standard error of the mean.
G.I. Details for Figure
The MDP considered is a two-state MDP, with single action a and discount factor O. Writing the states as x, y, we have P(2 Ir. a) 1 and Pyy y. = 1, that each state transitions itself. The reward distributions are given as R(x, a) = ōu. Ryy. a) { 10 + 510). Cramér-CTD and KL-CTD are both run using 40 atoms uniformly spaced on [-30, 30],
learning rate of 4 10-3 was used for TD and Cramér-CTD, and learning rate of 1 10-1 was used for KL-CTD.
G.2. Details for Figure 3
We two-state MDP with single action, deterministic rewards, and a discount factor of 0.7. The transition matrix is sampled from Dirichlet(1,.--- 1) distribution, and the reward means sampled independently from (0, KL-CTD was run for 5,000 iterations with 50 atoms uniformly spaced on [-10, 10] and a learning rate of 0.5.
G.3. Details for Figure 4
We use single-state MDP with single action and discount factor O. The return from this state is sampled from U[[-10, 10)). We use CTD with 10 atoms whose locations uniformly spaced across [-10, 10] with step size sequence given by Ok = k 00.55
G.4. Details for Figure 5
We consider an MDP with 10 states, transitions sampled from Dirichlet(1,.... 1) distribution, deterministic rewards whose means are sampled from U[0, 1J), and discount factor 0.5. We set the number of atoms to be 20, and uniformly space these atoms across [0, =m]. The final atom location 2m itself is swept over 100 uniformly spaced points in [0.1, 10].
G.5. Details for Figure 6
We consider an MDP with 20 states, transitions sampled from Dirichlet(1..... 1) distribution, deterministic rewards whose means are sampled from a standard normal (O, 1) distribution, and discount factor 0.9. We sweep over the number of atoms, m, in the set [4, 10, 16, 50, 100, 200]. For each choice of m, we set the locations to be uniformly spaced on [-10, 10J. We swept over learning rates uniformly-spaced in log-space over the range [1 10-4, 10]. For the scaled learning rates, we scale each learning rate by /m, so that learning rate of a corresponds an effective learning rate of a.m.
G.6. Details for Figure
All MDPs considered have 20 states and a discount factor of 0.7. All reward functions have their means sampled from a standard normal (0, 1) distribution. All methods are initialized so that their initial value estimates are O, and CTD and PCTD use 50 bins uniformly spaced [-100, 100). We sweep over learning rates uniformly-spaces in log-space over the range [10 04, 10). We now detail each type of transition and reward structure.
Transition dynamics
Cycle domain: We use deterministic cycle transition structure over the states.
Garnet domain: We sample sparse Gareet MDP transition structure (Archibald al., 1995).
Dirichlet domain: Each of the transition matrix is sampled i.i.d. from a Dirichlet(1,.... 1) distribution.
Reward distributions
Deterministic distributions:
The distributions are Dirac distributions centred at the means.
Gaussian distributions: The reward distributions are Gaussian with the given means with variance 1.
to distributions: The reward distributions are shifted to distributions to maintain the specified reward means.
27


Reviews:
[{"review_id": "S791oe01jP", "reviewer": "Reviewer_UawD", "summary": "The paper studies categorical distributional reinforcement learning  with a KL divergence loss. Unlike previous analyses relying on the Cramér distance, this paper introduces a novel preconditioned version of categorical temporal-difference learning with KL divergence , proving its convergence under mild assumptions. The paper also analyzes the asymptotic variance behavior of categorical estimates under various learning rate schedules. Empirical evaluations demonstrate that KL-based algorithms perform differently from classical temporal-difference methods in specific tabular reinforcement learning environments.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort in reviewing our paper, and for the helpful feedback provided. We are pleased that the reviewer found our work on the convergence of PKL-CTD, the asymptotic variance analysis, and the connection between theory and empirical observations to be valuable. Below, we address each of the reviewer’s points in detail.\n\n\n**Analysis is restricted to synchronous policy evaluation scenarios**\n\n*Synchronous/Asynchronous*\n\nWe note that all experimental results in Section 7 are done using asynchronous updates (with synchronous updates performed as an ablation in Appendix F), though the reviewer is correct that our convergence result addresses the synchronous setting. In light of the reviewer’s suggestions we have expanded our analysis, and give a sketch below of how it can be applied to analyse asynchronous cases too. \n\nFollowing work presenting core convergence guarantees in asynchronous stochastic approximation, such as Borkar (2008), considering the case where states are updated according to an ergodic Markov chain with stationary distribution $c$, the key concern to establishing convergence under appropriate technical conditions on step sizes is to establish a corresponding Lyapunov function for the continuous-time dynamics\n\n$$ \\partial_t \\phi_t (x) = c(x) C^\\top C (T^\\pi - I) p^{\\phi_t}(x),  $$\n\nwhich take into account the average per-state weighting of updates. From our analysis of the synchronous case (see Proposition 5.4), it can then be verified that $L(\\phi) = \\sum_{x\\in\\mathcal{X}} \\frac{d^\\pi(x)}{c(x)} \\mathrm{KL}(\\tilde{p}^\\pi(x) \\| p^\\phi(x))$ is a Lyapunov equation for this ODE, which can then be used to guarantee its convergence. \n\n*Policy Evaluation/Control*\n\nWe focus on the policy evaluation setting in this paper, and leave the control setting as interesting future work. We remark that understanding the policy evaluation setting first is typical in the analysis of both standard and distributional reinforcement learning.\n\n\n**Proof techniques used are relatively standard, existing stochastic approximation theory**\n\nWe agree that our proof of Theorem 5.5 uses existing stochastic approximation theory, although we would argue that many applications of stochastic approximation theory to RL make use of core existing theory (such as in the classic papers analyzing tabular TD algorithms by Jaakkola et al. (1994) and the linear TD analysis of Tsitsiklis and Van Roy (1997)), and indeed novel theory and techniques of stochastic approximation are often published as stochastic approximation results in their own right.\n\nFurther, we want to emphasize that our theoretical analysis leads to a number of novel results, such as an undiscovered algorithm (PKL-CTD), the effective learning rate phenomenon which we show is fundamental to KL-based categorical algorithms, and the bias-variance tradeoff in the atom locations.\n\n\n**Convergence of non-preconditioned KL-CTD**\n\nWe agree that the reviewer is entirely correct that our counterexample to the KL being a Lyapunov function for non-preconditioned KL-CTD does not exclude its possibility of convergence, and we will make this clearer in the paper. We chose to include this as an additional result in the appendix as it is a natural question in light of Proposition 5.4 (does the weighted KL, our exhibited Lyapunov function for PKL-CTD, also work as a Lyapunov function for KL-CTD?), we did not intend to suggest that non-preconditioned KL-CTD cannot converge.\n\n\n**Recent works on quantile-based DRL are not cited**\n\nWe would like to ensure that no relevant work is missing from our discussion of related work in Appendix E. We invite the reviewer to suggest any works in particular they find to be missing.\n\n\n**Other Comments Or Suggestions**\n\nWe thank the reviewer for all of the additional catches/suggestions, and we will ensure to fix them in the text.\n\n**References**\n\nJaakkola, Tommi, Michael Jordan, and Satinder Singh. \"Convergence of stochastic iterative dynamic programming algorithms.\" Neural Computation (1994).\n\nTsitsiklis, John and Van Roy, Ben \"An analysis of temporal-difference learning with function approximation,\" in IEEE Transactions on Automatic Control (1997)\n\nBorkar, Vivek. Stochastic approximation: a dynamical systems viewpoint. Cambridge University Press, (2008).", "claims_and_evidence": "The main claims inluding convergence of PKL-CTD, the asymptotic variance analysis, and advantages of KL-based methods are clearly supported by theoretical proofs and empirical experiments. However, the manuscript occasionally lacks clarity, especially around certain critical derivations (e.g., derivation of equation (6) from equation (5)). Additionally, the authors should more clearly differentiate their contributions from prior work. I do have some additional concerns\n- The paper provides a counterexample showing KL divergence does not serve as a Lyapunov function, but that does not exclude the possibility of convergence of KL-CTD without preconditioning.\n- The paper presents empirical results showing distinct learning dynamics for KL-CTD. However, the theoretical justification is somewhat heuristic, and the paper lacks a rigorous explanation for why KL-CTD is preferred beyond empirical observations.\n- The clarity regarding the derivation of certain critical equations (such as equation (6)) needs improvement.", "methods_and_evaluation_criteria": "The proposed methods (PKL-CTD and KL-CTD) and evaluation criteria, which uses standard RL benchmarks (Cycle, Garnet, Dirichlet environments), are reasonable and appropriate for assessing theoretical claims. However, the analysis is restricted to synchronous updates and policy evaluation settings without exploring more interesting setting like asynchronous or control (Q-learning), which limits the generality of the results.", "theoretical_claims": "The main convergence proofs appear correct and leverage standard stochastic approximation techniques and existing convergence theory. No technical flaws were found, though the techniques used are relatively standard and do not introduce significant novel technical complexities.", "experimental_designs_or_analyses": "The experimental analyses  illustrate the theoretical findings. Experiments comparing TD, KL-CTD, and PKL-CTD are well-designed and show that KL-CTD outperforms TD in high-stochasticity environments, while PKL-CTD is the most stable.", "relation_to_broader_scientific_literature": "The paper extends prior work on distributional RL (Bellemare et al., 2017) and categorical RL (Rowland et al., 2018). It is related to Cramér-based CDRL (Boeck & Heitzinger, 2022) but focuses on KL divergence instead. Some recent works on quantile-based DRL are not cited.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths\n- Theoretical advancement by addressing the convergence gap for KL-based distributional RL, which partially solved the problem raised by  Dabney et al., 2018\n- Asymptotic variance analysis is solid\n- Connects theory with empirical observations, particularly in variance behavior.\n- Derives insights on learning rate selection and category scaling\n\nWeakness\n-  Analysis is restricted to synchronous policy evaluation scenarios, not covering asynchronous and control setting\n- The proof of Theorem 5.5 is largely standard stochastic approximation theory, meaning the novelty lies more in the application rather than the mathematical difficulty.\n- Minor clarity issues in crucial derivations.", "other_comments_or_suggestions": "- Clarify the derivation from equation (5) to equation (6) explicitly in the main text or clearly reference its derivation in supplementary material.\n- Highlight the key differences/novelty from Rowland et al. (2018) and Boeck & Heitzinger (2022).\n- Provide more intuition for PKL-CTD by explaining why preconditioning improves convergence.\n- It might be insightful to discuss the potential impacts of different choices of the preconditioning matrix on convergence behavior and practical performance.\n- Indicates where the proof can be found \n- Proposition 4.3. The author should clarify that the stationary point is fixed point of Projectied Bellman operator\n- The equation in Line 224-228 is not clear\nSome typos\n- $|| ||_{\\pi}$ is not defined in Proposition 5.3.\n- Proposition 6.1. middle line $\\tilde{p}$", "questions_for_authors": "- Have you explored  the convergence and performance of KL-CTD or PKL-CTD with asynchronous updates and control setting?\n- While the paper provides a counterexample showing KL fails to be a Lyapunov function for KL-CTD, but that does not exclude the possibility of convergence of KL-CTD without preconditioning.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eEm5mctB5G", "reviewer": "Reviewer_wwL4", "summary": "This paper revisited categorical distributional RL and conducted the analysis based on KL divergence instead of the conventional Cramer distance. The authors proposed a variant of the algorithm by preconditioning and showed its convergence. More importantly, the asymptotic normality or asymptotic variance is also provided in estimating the values. Experiments are conducted in some simple environments to demonstrate the theoretical results.\n\n## update after rebuttal\n\nThank you for the response. This paper has many technical contributions, but I still think it could be largely enhanced. In particular, the authors can consider strengthening the motivation for why we need to fill the gap between Cramer CTD and KL CTD here, beyond only providing some simple motivating examples. The theoretical contributions should be made clearer. Also, I would suggest one or two conclusions be emphasized in the end instead of applying a tree-manner writing style that directly stacks multiple insights. Thus, I keep my rating.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort they spent in reviewing our paper, and we aim to resolve their questions below.\n\n**Why do the authors consider the KL variant of categorical distributional RL, and what is the gap between it from the original algorithm, such as C51?**\n\nWe focus on the KL variant of categorical distributional RL as it is exactly the fundamental algorithm used by C51. The gap between KL-CTD considered here and C51 is due to the use of policy evaluation instead of control, and tabular analysis instead of the deep learning setting.\n\n**The motivation for the study of KL-CTD is not clear in Section 3. Why do we not directly study the original algorithm?**\n\nWe motivate KL-CTD in two ways in Section 3:\n\n- Firstly, KL-CTD is the original categorical TD algorithm that C51 makes use of. This in itself is a strong motivation for its study.\n- Our examples in Figure 2 and 3 serve as further motivation, showing that even in tabular settings, without neural network function approximation, KL-CTD can exhibit drastically different behaviour to earlier studied algorithms such as Cramer CTD and classical TD learning.\n\nWe emphasize that KL-CTD is the original CTD algorithm, as described at Eqn (6) and Line 141.\n\n\n\n**How to understand the convergence result in this paper based on KL divergence, given that distributional dynamic programming is only non-expansive if we directly employ the KL divergence to measure the distribution distance**\n\nOur convergence analysis relies on Lyapunov stability results, as opposed to writing our updates as a stochastic approximation of a contractive operator. \n\nMore specifically, there are several reasons why a result concerning the unprojected distributional Bellman operator's behaviour as measured by KL divergence is not pertinent to our analysis here:\nWe are concerned with convergence of an algorithm that maintains a categorical representation of the return distribution, but the result above is concerned with the unprojected distributional Bellman operator.\nWhile we analyze an algorithm that is defined via the KL loss, this does not obligate us to make use of contraction results measuring distance in KL. In Bellemare et al. (2023, Chapter 5), this difference is emphasised by distinguishing between metrics used for analysis, and metrics used to define algorithms. In our case, we study an algorithm that uses KL between categorical distributions in its definition, and our analysis makes use of both KL (to measure distance from the fixed point), and Cramér distance (we make use of contractivity of the categorical distributional Bellman operator in Cramér distance to show that the KL is a Lyapunov function; see the proof of Proposition 5.4 and line 783 in particular).\n\nPlease let us know if you have any further queries on this point.\n\n\n**$\\tilde{p}$ above Equation (8)**\n\nThe $\\tilde{p}^\\pi_i$ appearing in the equation $\\tilde{\\eta}^\\pi = \\sum_i \\tilde{p}^\\pi_i \\delta_{z_i}$ represents the probabilities of the fixed point of the projected distributional Bellman operator.\n\n​​\n**It is not clear why C is introduced in Eq.9 and 10. Could any intuitive explanation be provided?**\n\nOur previous calculation on lines 225-227 shows that the KL divergence from the current estimate to the fixed point may be increasing under the KL-CTD dynamics. But from this expression, in light of Proposition 5.3., we can reverse-engineer a modification to the updates which would guarantee the decrease of the KL, and lead to convergence. This modification is exactly the preconditioning using $C^\\top C$.\n\n\n**It seems that the superiority of the proposed algorithms PKL-CTD over KL-CTD and TD highly depends on the step size.**\n\nThe reviewer is correct that which algorithm is best between the various considered in Figure 7 depends on the step size, however we would like to emphasize that at least in the settings in the figure, the dependence of PKL-CTD on step size has a generally equal or smoother curve than the dependence of KL-CTD and TD on step size, which both have deep learning equivalents which are highly used in practice (DQN and C51).\n\n\n**Why consider the asymptotic normality/variance in value estimates in Propositions 6.1 and 6.4? What is the main regularity condition? Is there any potential application of these asymptotic properties?**\n\nWe study the asymptotic distributions of the value estimates because they are able to give us exact forms of the errors incurred by the algorithms. These asymptotic results have led to the bias/variance results in our paper, and we believe that they have further applicability for deriving insights in future work.\n\n\n**Comparison with Rowland (2018)**\n\n\nWhile Rowland et al. (2018) study categorical dynamic programming and the Cramer-CTD update, this paper focuses on the categorical TD learning algorithm based on the gradient of the KL loss, matching the form of loss used in the original categorical distributional RL paper (Bellemare et al., 2017).", "claims_and_evidence": "This paper gives me the feeling that many parts are disconnected without presenting one or two clear conclusions, even though each part is involved with rigorous analysis. For example, why do the authors consider the KL variant categorical distributional RL, and what is the gap between it from the original algorithm, such as C51? Why do we then move to a preconditioned variant of an algorithm rather than rethink the properties of real categorial distributional RL algorithms? Furthermore, I understand it takes effort to conduct the convergence and asymptotic analysis in value estimates, but what are the purposes or goals of presenting these results? Do the authors hope to demonstrate an advanced algorithm that can be superior to the original categorical distributional RL or just ``dump’ ’related properties that could be derived? After reading this paper, I am very confused about the main purpose of this paper and what kind of insights or conclusion the readers should gain after the reading.", "methods_and_evaluation_criteria": "What is the motivation to consider categorical distributional RL with KL divergence? Is that because the previous analysis in (Rowland et al. 2018), such as Cramer TD, does not align well with the real practical algorithm? Or is it just because the KL variant algorithm could simplify the theoretical analysis?\n\nThe motivation for the study of KL-CTD is not clear in Section 3, given that the environments are particularly chosen. In addition, the original categorical distributional RL may also behave erratically in Figure 3. Why do we not directly study the original algorithm?\n\nI find it difficult to understand the motivation of preconditioning variants. It seems to be mainly from a theoretical perspective.", "theoretical_claims": "The theoretical claims look fine to me. One question is how to understand the convergence result in this paper based on KL divergence, given that distributional dynamic programming is only non-expansive if we directly employ the KL divergence to measure the distribution distance.\n\nIs $\\tilde{p}_i^{\\pi}$ above Eq. 8 in $\\tilde{\\eta}$ the probability after the two-hot transformation?\n\nDoes Proposition 5.3 mean that KL-CTD converges to the same fixed point of the original categorical distributional RL algorithm?\n\nIt is not clear why $C$ is introduced in Eq.9 and 10. Could any intuitive explanation be provided?\n\nWhy consider the asymptotic normality/variance in value estimates in Propositions 6.1 and 6.4? What is the main regularity condition? Is there any potential application of these asymptotic properties?", "experimental_designs_or_analyses": "Although this is a theoretical paper, the experiments are still weak from my perspective. For example, Figures 2 and 3 are on a particularly chosen environment, which is less convincing. The results in Figure 4 are appealing, but they are conducted on a very simple one-state environment after I checked the details in the appendix. In addition, in Section 7.2, it seems that the superiority of the proposed algorithms PKL-CTD over KL-CTD and TD highly depends on the step size (learning rate), which is frustrating, especially for practitioners. It is suggested to have experiments on some classical gym control environments, even though it may not be necessary to consider complex environments for a theory-focused paper.", "relation_to_broader_scientific_literature": "Although I agree that the theoretical analysis looks rigorous for this paper, my concern is that it may not be related to the broader scientific literature. The focus of the paper is a modification of categorical distributional RL, which is one specific distributional RL algorithm. Although this modification with KL divergence simplifies the theoretical analysis, allowing rich convergence and asymptotic normality analysis in value estimates, it may still be specific. I am also looking for some potential theoretical contributions that may be generalized to broader areas. However, I feel that many analyses are based on existing results and are established in a parallel manner. Therefore, I am concerned about whether this paper could be related to broader literature, given that the analytic tools may be a little specific in a variant of categorical distributional RL.", "essential_references_not_discussed": "I think the references are relatively self-contained. It is a theory-focused paper that mainly discusses the related work directly linked with the analysis, but it would be better to add more references on the whole distributional RL area.", "other_strengths_and_weaknesses": "Please see the comments above.", "other_comments_or_suggestions": "What is the gap between KL-based categorical distributional RL and the vanilla algorithm?", "questions_for_authors": "It looks like this is a follow-up paper to Rowland (2018). What is this paper's main novelty or contribution relative to it?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nmsg9BBMPg", "reviewer": "Reviewer_YLKh", "summary": "In this paper, the authors studied the theoretical properties of categorical distributional TD with KL loss. They proposed a preconditioned version of the algorithm called PKL-CTD, proved its asymptotic convergence, and derived the asymptotic distribution of the resulting value estimators. These theoretical results also provide valuable insights for practitioners.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and effort in reviewing our paper, and for the helpful feedback and suggestions provided. We are pleased to hear that they found our work to provide valuable insights for practitioners and our idea of preconditioning to be concise and intuitive.\n\n**Asymptotic analysis missing for Cramér-CTD**\n\nWe thank the reviewer for pointing out this limitation, and we have expanded the asymptotic results based on their feedback. In particular, we have derived additional results for the asymptotic normality of $p_k^{Cramer-CMC}$, and present this below.\n\nLet us define the iterates $p_k^{\\text{Cram\\’er-CMC}}$ as follows: $p_{k+1}^{\\text{Cram\\’er-CMC}} = p_k^{\\text{Cram\\’er-CMC}} + \\alpha_k (h(G) - p_k^{\\text{Cram\\’er-CMC}})$. Then under the assumptions of Lemma D.3., we have that\n\n$$\nk^\\beta ( p_{k}^{\\text{Cram\\’er-CMC}}  - \\tilde{p}^\\pi)  \\overset{d}{\\to} \\mathcal{N}\\left(0, \\mathbb{E}_{G}[h(G) h(G)^T] - (\\tilde{p}^\\pi)(\\tilde{p}^\\pi)^T\\right).\n$$\n\n\n**Extension of asymptotic analysis to general statistical functionals** \n\nWe thank the reviewer for highlighting this as a direction to expand our current results, as we find this to be a nice generalization of our current analysis. We present the result for KL-CMC below, and will include it for TD and Cramér-CMC in the paper.\n\nSuppose $\\psi: \\mathscr{P}(\\mathbb{R}) \\to \\mathbb{R}^k$ is a statistical functional sketch, and let $J_\\psi: \\mathbb{R}^m \\to \\mathbb{R^k}$ be the Jacobian of $\\phi \\mapsto \\psi(p^\\phi)$. Then if $J_\\psi$ is continuous in a neighbourhood of $\\tilde{p}^\\pi$, we have\n\n$$\nk^\\beta ( \\psi(p_{k}^{\\text{KL-CMC}})  - \\psi(\\tilde{p}^\\pi))  \\overset{d}{\\to} \\mathcal{N}\\left(0, \\frac12 J_\\psi(\\tilde{p}^\\pi) \\left(\\mathrm{diag}(\\tilde{p}^\\pi) - (\\tilde{p}^\\pi)(\\tilde{p}^\\pi)^\\top\\right) J_\\psi(\\tilde{p}^\\pi)^\\top\\right).\n$$\n\n\n**Work only considers the tabular case**\n\nWe appreciate that while our focus on the tabular setting potentially seems a limitation due to its simplicity, we want to emphasize that this was a conscious choice in order to develop a core understanding of categorical TD algorithms based on KL gradients, which have not been studied before, without confounding factors that often arise in larger-scale experiments (such as adaptive optimisation, function approximation, target networks, etc.).\n\n\n**Missing blue line in Figure 2**\n\nThe line for the MSE of TD is not missing, it rather perfectly overlaps with the line of Cramér-CTD, as predicted by the theory of Lyle et al. (2019). This is a key motivation for the work: while the behaviour of Cramer-CTD exactly coincides with classical TD learning, as predicted by Lyle et al. (2019), this shows that there are settings with KL-CTD has distinct, and sometimes superior, performance to these other algorithms, motivating us to understand the theory of this approach better. We will make this more emphasized in the text, and explore other methods of visualization to remove any confusion, such as making the linewidth of the blue curve larger.\n\n\n**References**\n\nLyle, C., Bellemare, M. G., and Castro, P. S. A comparative analysis of expected and distributional reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2019.", "claims_and_evidence": "The claims made in the submission are supported by sound theoretical analysis or empirical evidence.", "methods_and_evaluation_criteria": "The proposed methods make sense for the problem at hand.", "theoretical_claims": "The theoretical results seem reasonable in this work. I checked most parts of the proof and found them correct.", "experimental_designs_or_analyses": "The experiments are adequate for the purpose of validating the theoretical findings.", "relation_to_broader_scientific_literature": "This paper presents a thorough discussion of the theoretical properties of categorical distributional TD learning with KL loss. The theoretical findings advance the understanding of the distributional TD learning algorithm and provide valuable insights for practitioners. I think the authors has made a solid contribution to the field of distributional RL.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n\n* The idea of introducing a preconditioner is concise and intuitive.\n* The asymptotic analysis of the value estimator provides a series of valuable insights, for example, the bias-variance tradeoff incurred by the choice of the size of supports.\n\nWeaknesses:\n\n* The asymptotic analysis only focuses on the asymptotic behavior of value estimators in TD and KL-CTD. I think it may also be helpful to compare the asymptotics between Cramer-CTD and KL-CTD.\n* As the authors state, \"However, most large-scale implementations of categorical temporal-difference learning use a KL loss, rather\nthan Cramer loss. This is a crucial detail of large-scale implementations, but has not yet been theoretically analyzed.\" However, it seems that the author only considers the tabular case. I think the authors should at least use some experiments to show that the insights obtained in the tabular case are also valid in the large-scale case with the function approximation techniques applied.", "other_comments_or_suggestions": "I do not see the blue line indicating the MSE of TD learning in Figure 2, I hope the author can fix it in later revisions.", "questions_for_authors": "Can the authors generalize the asymptotic analysis to estimators of statistical functionals other than the mean? For example, the asymptotic distribution of differentiable statistical functionals can be obtained with a combination of asymptotics of the learned weights and delta method.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sj2zdW2bwK", "reviewer": "Reviewer_tUof", "summary": "This paper analyzed categorical TD learning with KL loss in the tabular setting. They also proposed a preconditioned version of the algorithm and provided an asymptotic normal analysis. They also conducted experiments to verify the theoretical results.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time and effort in reviewing our paper, and we are pleased that they found our motivation compelling, appreciated our use of preconditioning to obtain a convergence guarantee, and found our use of asymptotic variance for the analysis to be a clear choice.\n\n\n**Comparison between Cramér and KL losses for learning distributions**\n\nWe thank the reviewer for this suggestion, as we agree that an asymptotic comparison of the probability vectors learnt by the Cramér and KL losses is a valuable comparison to make. We report this result in a new proposition, which we state below.\nLet us define the iterates $p_k^{\\text{Cram\\’er-CMC}}$ as follows: $p_{k+1}^{\\text{Cram\\’er-CMC}} = p_k^{\\text{Cram\\’er-CMC}} + \\alpha_k (h(G) - p_k^{\\text{Cram\\’er-CMC}})$. Then under the assumptions of Lemma D.3., we have that\n\n$$k^\\beta ( p_{k}^{\\text{Cram\\’er-CMC}}  - \\tilde{p}^\\pi)  \\overset{d}{\\to} \\mathcal{N}\\left(0, \\mathbb{E}_{G}[h(G) h(G)^T] - (\\tilde{p}^\\pi)(\\tilde{p}^\\pi)^T\\right).$$\n\n\nWe also want to emphasize that the goal of our work is not necessarily to argue that the KL version of the algorithm is always superior to the Cramer-based version (and indeed this is not always the case).\nOur main motivation is to study the KL version in its own right, which is the fundamental algorithm underlying the original deep RL implementation of distributional RL (C51; Bellemare et al., 2017).\n\n\n\n\n\n**Experimental comparisons between the KL loss and the Cramer loss under the function approximation setting**\n\nWe appreciate that while our focus on the tabular setting potentially seems a limitation due to its simplicity, we want to emphasize that this was a conscious choice in order to develop a core understanding of categorical TD algorithms based on KL gradients, which have not been studied before, without confounding factors that often arise in larger-scale experiments (such as adaptive optimisation, function approximation, target networks, etc.).\n\nWe additionally note that the paper suggested by the reviewer performs an empirical comparison of KL-CTD and Cramér-CTD in the function approximation setting: C51 is the deep learning equivalent of KL-CTD, and S51 as used in their paper is the deep learning equivalent of Cramér-CTD.\n\n\n**Other Comments Or Suggestions**\n\nWe thank the reviewer for bringing attention to these typos. We will be sure to fix them in the text.\n\n\n**References**\n\nBellemare, M. G., Dabney, W., and Munos, R. A distributional perspective on reinforcement learning. In International Conference on Machine Learning, 2017.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes", "relation_to_broader_scientific_literature": "This paper analyzes categorical TD learning with KL loss, which are used in many deep distributional reinforcement learning algorithms.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strength:\n1. The motivation is good. On the one hand, the KL loss (instead of the Cramer loss) is adopted by many practical algorithms, and studying the performance under this loss function helps in understanding real-world algorithms. On the other hand, it also inspires the development of distributional reinforcement learning algorithms that utilize other probabilistic metrics as loss functions. \n2. It is a good idea to use preconditioning to improve KL-CTD and guarantee convergence.\n3. For readers, it is quite clear to use the asymptotic variance to compare the performance of algorithms in the simple Monte-Carlo and tabular settings.\n\nWeakness:\n1. Although the motivation is quite good, and there are also experimental and some theoretical evidences to illustrate the advantages of using the KL loss compared with the Cramer loss, considering that this paper focuses on the simplified tabular setting, I think it would be better to have theoretical evidence demonstrating that, in terms of learning distributions (not just the means, i.e. value functions in Section 6), the KL loss has advantages over the Cramer loss. For example, you can provide the asymptotic normality of the vector p_k^{Cramer-CMC} and p_k^{KL-CMC}  \n2. The paper mentioned that the KL loss is used in many practical algorithms. If there could be experimental comparisons between the KL loss and the Cramer loss (for example, see linear function approximation algorithm in [1]) under the function approximation setting, I believe the paper would be more complete.\n[1] Lyle, C., Bellemare, M. G., and Castro, P. S. A comparative analysis of expected and distributional reinforcement\nlearning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2019.", "other_comments_or_suggestions": "Some notations are used without being defined., such as \\|... \\|_{\\pi}, d^\\pi.\nThere is typo in the equation in Line 225-228.", "questions_for_authors": "No", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper analyses distributional temporal difference learning with categorical representation using KL divergence as loss (KL-CTD). The choice of KL divergence matches the original C51 paper while previous analysis relies on changing the KL divergence to some other loss to facilitate analysis. Based on this fact, I believe this paper is well motivated. That being said, as pointed by the reviewers, this paper does fall short in terms of experiments and considering only synchronous updates in the analysis. I personally believe the extension to asynchronous updates should be simple, especially with i.i.d. samples. Missing this extension does not constitute a major flaw to me. And I appreciate the authors' initial efforts towards this in the response. Since the contribution is mostly theoretical, I will not hold a very high standard for empirical study to require large scale experiments. Two reviewers are negative but I do not see their complain about technical flaws. 

Overall, I think this is a sold work that nicely closes a gap in the literature. I therefore recommend accept. That being said, I do encourage the authors to improve the paper according to the reviews.
Decision: Accept (poster)