One-Step Distributional Reinforcement Learning
Authors: Mastane Achab
Reda ALAMI
YASSER ABDELAZIZ DAHOU DJILALI
Kirill Fedyanin
Eric Moulines
Venue: Accepted by TMLR ()

Abstract: Reinforcement learning (RL) allows an agent interacting sequentially with an environment to maximize its long-term expected return. In the distributional RL (DistrRL) paradigm, the agent goes beyond the limit of the expected value, to capture the underlying probability distribution of the return across all time steps. The set of DistrRL algorithms has led to improved empirical performance. Nevertheless, the theory of DistrRL is still not fully understood, especially in the control case. In this paper, we present the simpler one-step distributional reinforcement learning (OS-DistrRL) framework encompassing only the randomness induced by the one-step dynamics of the environment. Contrary to DistrRL, we show that our approach comes with a unified theory for both policy evaluation and control. Indeed, we propose two OS-DistrRL algorithms for which we provide an almost sure convergence analysis. The proposed approach compares favorably with categorical DistrRL on various environments.

Published in Transactions on Machine Learning Research (08/2023)
One-Step Distributional Reinforcement Learning
Mastane Achab
mnsstone.cchtee Atäi.
Reda Alami, Yasser Abdelaziz Dahou Djilali, Kirill Fedyanin Technology Innovation Institute, 9639 Masdar City, Abu Dhabi, United Arab Emérates
Eric Moulines
Ecole polytechnique, Palaiseau, France
Reviewed OpenReview: https: https://openrevieu openreview. net/ forum???? id= ZPMf53uE1L
Abstract
Reinforcement learning (RL) allows agent interacting sequentially with an environment to maximize its long-term term expected return. In the distributional RL (DistrRL) paradigm, the agent goes beyond the limit of the expected value, to capture the underlying probabil- ity distribution of the return across all time steps. The set of DistrRL algorithms has led to improved empirical performance. Nevertheless, the theory of DistrRL is still not fully understood, especially in the control case. In this paper, we present the simpler one-step distributional reinforcement learning (OS-DistrRI) framework encompassing only the ran- domness induced by the one-slep dynamicc of the environment. Contrary to DistrRL, we show that our approach comes with a unified theory for both policy evaluation and control. Indeed, we propose two OS-DistrRL algorithms for which we provide an almost sure con- vergence analysis. The proposed approach compares favorably with categorical DistrRL on various enviromments.
1 Introduction
In reinforcement learning (RL), decision-maker agent sequentially interacts with unknown and uncertain environment in order to optimize some performance criterion (Sutton & Barto 2018). At cach time step, the agent observes the current stale of the environment, then takes an action that influences both its immediate reward and the next state. In other words, RL refers to learning through trial-and-error a stratogy (or policy) mapping states to actions that maximizes the long-term cumulative reward (or return): this the so-called control lask. More accurately, this return a random variable due to the random transitions across states and classical RT only focuses on its expected value. On the other hand, the policy evaluation task aims assessing the quality of any given policy (not necessarily optimal as in control) by computing its expected return in each initial state, also called value function. For both evaluation and control, when the model (i.e. reward function and transition probabilities between states) is known, these value functions can be seen as fixed points of some operators and computed by dynamic programming (DP) under the Markov decision process (MDP) formalism (Puterman 2014) Nevertheless, the model in RL is typically unknown and the agent can only approximate the DP approach based on empirical trajectories. The (D( algorithm for policy evaluation and Q-learning for control, respectively introduced in (Sutton 1988)
and (Watkins & Dayan 1992), are flagship examples of the RL paradigm. Formal convergence guarantees were provided for both of these methods, see (Dayan 1992 Tsitsiklis. 1994 Jaakkola et al.) 1993). In many RL applications, the number of states is very large and thus prevents the use of the aforementioned tabular RL algorithms. In such situation, one should rather use function approximation to approximate the value functions, as achieved by the DQN algorithm (Mnih et al., 2013 2015) combining ideas from Q-learning and deep learning. More recently, the distributional reinforcement learning (DistrRL) framework was proposed by Bellemare et al. (2017); see also (Morimura al.) 2010a b). In DistrRI, the agent is optimized to model the whole probability distribution of the return, not just its expectation. In this new paradigm, several distributional procedures where proposed extensions of the classic (non-distributiionl)

---

Published in Transactions on Machine Learning Research (08/2023)
RL methods, leading to improved empirical performance. In most cases, DistrRL algorithm is composed of two main ingredients: i) a parametric family of distributions serving as proxies for the distributional returns and ii) a metric measuring approximation errors between original distributions and parametric proxies.
Related work. We recall a few existing DistrRL, approaches, all of which considering mixtures of Dirac measures as their parametric family of proxy distributions. Indeed, the Categorical DistrRL (CDRL) parametrization used in the C51 algorithm proposed in (Bellemare al., 2017 was shown in (Rowland) et al. 2018) to correspond to orthogonal projections derived from the Cramér distance The CDRL ap- proach learns the categorical probabilities pı (I, a..... ... Pk (I, a) a distribution Σx= (t, a)o2, with fixed predefined support values 21,..., =k. On the other hand, the quantile regression approach was proposed in (Dabney et al. 2018), where the support {Qi (z..)..... ... Qx a}} of the proxy distribution 1 ( is learned for fixed uniform probabilities N over the N atoms. In (Dabney et al. 2018), these atoms result from Wasserstein-1 projections, and correspond to quantiles of the unprojected distribution. [Achab & Neu 22021 and Achab et al. (2022) later investigated the Wasserstein-22 sctup leading to the study of conditional value-at-risk measures. For policy evaluation, the convergence analysis of CDRL and of the quantile approach were respectively derived in (Rowland et al. 2018) and (Rowland et al.
2023) Nevertheless, the theory of DistrRL is more challenging in the control case because the corresponding operator is not a contraction 2 and does not necessarily admits a fixed point. For that reason, Rowland et al. (2018) provod the convergence of CDRL for control only under the restrictive assumption that the optimal policy is unique. Hence, it seems natural to ask the following question: "Is there another formulation of DistrRL in which both the evaluation and the control tasks lead to contractive operators?". The answer provided by this paper is "Yes, via a one-step approach!". Contrary to classic DistrRL dealing with the randomness across all time steps, the one-step variant (originally proposed by Achab 2020b) only cares about the one-step dynamics of the environment.
Contributions.
Our main contributions are as follows:
We introduce one-step variant of the DistrRL framework: call that approach one-step distri- butional reinforcement learning (OS-DistrRL).
We show that our new method solves the well-known instability issue of DistrRL in the control case.
We provide a unified almost-sure convergence analysis for both evaluation and control.
We experimentally show the competitive performance of our new (deep learning-enhance)) algo- rithms in various environments.
The paper is organized as follows. In Section 2 we recall a few standard RL tools and notations as well as their DistrRL generalization. Then, our one-step approach is defined in Section g Section 4 introducess new OS-DistrRI. algorithms along with theoretical convergence guarantees. Finally, numerical experiments are provided for illustration purpose in Section The main proofs are deferred to the Supplementary Material.
Notations. The indicator function of any event F. is denoted by T(... We let Po(R) be the set of probability measures on I having bounded support, and P(E) the set of probability mass functions on any finite set ε, whose cardinality is denoted by [E). The support of any discrete distribution E P(E) is s4pport(4) = (y E E: q(y)) > 0); the supremum norm of any function h: E R is ||Allc = max [h(y))- The cumulative distribution function (CDF) of a probability measure v E P&(R) is the mapping F(z) = (Z S 2) (Vz E R), and we denote its generalized inverse distribution function (a.kaa quantile function) by F-1: (0, 1] + inf(z E R. r}. Given 21, 1/2 E (R) with respective CDFs Fi, F2, we denote 12 in and say that in stochastically dominates V2 if F1(2) < F2(2) for all R. For any probability
The Cramér distance l between two probability distributions v1, with CDFs Fi, equal 1a(2),1) =
(Z)(x) [a(x))2de
2 function mapping metric space itself called a y-contraction if Lipschitz continnous with Lipschitz constant.
7 < 1.
2

---

Published in Transactions on Machine Learning Research (08/2023)
measure (R) and measurable function f: R + R, the pushforward measure fgv' is defined for any Borel set A _ IR by (+)( (A)= m(
R: f(z) E AJ). In this article, we only need the affine case
(2))(2) ro + 72 (with ro E y 1)) for which # E (PR)) and f = u1 - 2 E A}) if
Y O, or f // 5.p is the Dirac measure at ro y =
2
Background on distributional reinforcement learning
In this section, we recall some standard notations, lools and algorithus used in RL and DistrRL.
2.1 Markov decision process
Throughout the paper, we consider a Markov decision process (MDP) characterized by the tuple P.r.7) with finite state space X, finite action space A, transition kernel P: X A + P(X), re- ward function r: X A x X + R, and discount factor 0 S Y < 1. If the agent takes while the enviromment is in state E X, then the next state X1 is sampled from the distribution P(- a) and the immediate reward is equal to r(z, X1). In the discounted MDP setting, the agent seeks policy w: X + P(A) maximizing its expected long-term return for each pair () of initial state and action:
action E A
Q"(z,a)= E
Σ 1=
Xx,, Xe+1) Xo = x. Ao =
where X1+1 P(IX,, Ad) and A++1 x(-1X++1). (z) and (=xx)= =
(aaxx) (x, a) are respec-
tively called the state-action value function and the value function of the policy Further, cach of these functions can be seen as the unique fixed point of a so-called Bellman operator (Bellman 1966). that we denote by 7 for Q-functions. For any Q: X A R. the image of Q by T is another Q-function given by:
+y Σ
This Bellman operator T has several nice properties: in particular, it is a 7-contraction in || ll and thus admits unique fixed point (by Banach's fixed point theorem), namely QT = T=Q". It is also well-known from (Bellman 1966) that there always exists at least one policy x" that is optimal uniformly for all initial conditions (r:
(x( = (x, = sup (x, and V*(x) = max (x, V= (*) = sup V*()).
a
Similarly, this optimal Q-function Q' is the unique fixed point of some operator T called the Bellman optimality operator and defined by:
(TQ) a) Σ () + max Q( ((d),a)))
which is also a r-cootraction lloo Noteworthy, knowing Q* sufficient to behave optimally: indeed, a policy x² is optimal if and only if every state x, support (x= (+((=)) arg max, Q" (x,a).
2.2 The distributional Bellman operator
In distributional RI, we replace scalar-valued functions Q by functions μ taking values that are entire probability distributions: (=2,a) e (R)) for each pair (zx, a). In other words, μ is a collection of distributions indexed by states and actions. We recall below the definition of the distributional Bellman operator, which generalizes the Bellman operator to distributions.
Definition 2.1 (DIstRIBUTIONAL BELLMAN OPERATOR, BELLEMARE ET AL. (017 Let æ be a policy. The distributional Bellman operator Th PA(RRRAAA - Po(R)AX is defined for any distribution function == (p(z,a)) ) by (T (T=n)(z,))= Σ (aal
3

---

Published in Transactions on Machine Learning Research (08/2023)
ax1
a,x)
a(x))
P(x2)x2,a)==1
a(x))
State X1
State X2
P(x1)x2.a,)=0
a(x))x2,a2)=055
P(x))22,22=0015
ax)
a2x
Figure Markov decision process with two states X (21, two actions A = [2) and reward function independent of the next state: r(z, a. r(z,aa) a).
We know from Bellemare al. (017)) that T is a 7-contraction in the maximal y- p-Wasserstein metrid W H2)
( (= 122 (1,a))
max (x.0)XXA
any order p +oo]. Consequently, T" has unique fixed point He (Hi ( )aa equal to the collection
of the probability distributions of the returns:
(a,0) Distr Σ AX,, Xx-1) No = x, Ao =
Applying 7 may be prohibitive in terms of space complexity: if μ is discrete, then 7 is still discrete but with up to |28 [A] times more atoms, shown in the following example. Example 2.2. Let μ (p(aa)) be the collection of the atomic distributions
K
u(a.0) Σ
pk la a) a)
where K 2 1, pe(x,, a) 2 and (1(x, a) + ..... + PK(1, = 1. Then, T is also a collection of atomic distributions with (at most) |x| [A] times more atoms:
(T (1) (a.a)
Σ P(a'(o,a)x(aal Σ ph (d..,..))rrraae)))))0eae
where we stress that the entire expression [r(æ, a, x') Qu d')
the argument of the Dirac delta function.
Projected operators. Motivated by this space complexity issue, projected DistrRL operators were pro- posed to ensure predefined and fixed space complexity budget. A projected DistrRL operator is the position of DistrRL operator with projection over parametric family of distributions. The Cramér distance projection Ie has been considered in (Bellemare et al. 2017 Rowland et al. 2018 Bellemare et al. 2019): we recall its definition below.
Definition 2.3 (CRAMÉR PROJECTION). Let K 2 2 and 21 <... < zÆ be real numbers defining the support of the categorical distributions. The Cramér projection is then defined by: for all z' E R,
8 if z' z' 21 B1=1=='
nc(62:)
z' 82,
+ 2'-2, 2'221 83 21-1-2z
if 2g < $ 2/+1
2541
85y if L 2K
3For p 2 1, recall that the Wasserstein distance between probability distributions V1 Uy on with CDFs Fi, Fg is defined as Wp(p1,02) 02) = (I'' ( (-) dr =fp co, (oo 12) = supteto.1) (T))(r) Fr)(

---

Published in Transactions on Machine Learning Research (08/2023)
0.6
L V W0
3.4
as.. ail (CDRL)
P2
0.6
3.2
3.0
pale..on
0.6
Pelks m1)
0.3
U VUU
0.4
2.6
2.4
2.2
0.0
0.0
2.0
20
40
60
80
100
20
40
60
100
20
40
60
80
100
(a) CDRL at (2t1, 01).
(b) One-step (21, 01).
(c) Both Q-functions at (11.01).
PL
Qit.. eg) (CDRL)
0.8
Pal
Oy), (One-step)
PalXt, 02)
3.0
84
2.8
03 WW
2.8
92-
0.2
0.0-
2.0
100
20
100
(d) CDRL (21, 02).
(e) One-step (x), 02).
(f) Both Q-functions (z1,a2).
Figure 2: Instability of categorical distributional control compared to the convergence of our one-step ap- proach. The two leftmost plots the iterations of the Cramér-projecte "full" optimality operator IlcT, while the center plots are obtained with our one-step operator IeT. The rightmost plots correspond to the Q-functions Qxx, a) Ei= Ph (I, a)
and extended linearly to mixtures of Dirac measures.
We will also apply II. entrywise to collections of distributions: if f (1(t,,)) then 1lcи (I)) The Cramér projection satisfies an important mean-preserving property: for any discrete distribution q with support included in the interval =x], then:
Ez~lle(g) [2] [Z]
(2)
In this line of work, the proxy distributions Σ -1 Pa (x, a))=, parametrized by the categorical probabilities Pi (z, (2,(x,u) In (Dabney et al. 2018 the Wasserstein-1 projection Ilw, over atomic distributions N D==1&0, is used. For this specific projection, the atoms that best approximate some CDF Fza are obtained through quantile regression with (x, = Frdd(33 21-1
2.3 Instability of distributional control
Bellemare et al. 2017 have also proposed DistrRL optimality operators T, defined such that Tu= T where NG is greedy policy with respect to the expectation of μ.. The authors showed in their Propositions 1-2 that, unfortunately, T is not a contraction and does not necessarily admits a fixed point. For that reason, the convergence analysis of the control case in DistRL more challenging than for the evaluation task. Rowland et al. (2018) and Bellemare et al. (2023) (sce chapter 7.4, in particular Theorem 7.9) circumvent this issue reducing the control task evaluation under the assumption that the optimal policy is unique. We start our investigation by verifying that this uniqueness hypothesis is critical ensure convergence in DistrRL control. For that purpose, we choose toy example fully described in Section 5 with (infinitely) many optimal policies and we run CDRL procedure over it. As expected, we observe in Figure 2 the unstable behavior of the DistrRL paradigm for the control task. Indeed, as explained in chapter 7.5 in Bellemare et al. (2023), multiple optimal policies can produce different distributions that are then mixed by DistrRL control algorithms which might never converge. In the next soction, we introduce our new "onc-step" distributional approach converging even in this situation since all optimal policies will be shown share the same fixed point distribution.
5

---

Published in Transactions on Machine Learning Research (08/2023)
3 One-step distributional operators
This section introduces the building blocks of our new framework. The main goal here is to show that our one-step approach is theoretically sound, though being less ambilious than "full" DistRLL. The main benefit of this simplification is that it will allow us to derive in the next section a convergence result holding for both evaluation and control. In particular in the control case, we will not need any additional assumption such as the uniqueness of the optinal policy as CDRL (Rowland et al. 2018
3.1 Formal definitions
Let us now define the one-step DistrRL operators. Intuitively, they are similar to the "full" DistrRL operator T" except that they average the randomness after the first random transition and are thus oblivious to the randomness induced by the remaining time steps.
Definition 3.1 (ONE-STEP DISTRIBUTIONAL BELLMAN OPERATOR). Let be a policy. The one-step Po(R) is defined for any distribution function by
distributional Bellman operator J=
(T p) (.,a)
Σ P(x a) ((,,,)))) E. =(a'(a)) 2.nptt.a)11)
where we stress that the argument of the Dirac delta function is: r(z,a,x)) + (a (azl pu ula' a [Z].
Definition 3.2 (ONE-STEP DISTRIBUTIONAL BELLMAN OPTIMALITY OPERATOR). The one-step distribu- tional Bellman optimality operator J: (R + (x) is defined for any μ by
(Tn))
-- P(z maax, E 2nnlt"
Example 3.3. Let j be the same discrete distribution function in Example 2.2
Then,
(T"))(z,a) = P(a' a) (( Eno ala (((()aa))))
and (Jn)(ana)= = Σ P(x (z, (1)0 Q(x',a')
k (= d')(x(z',d')) where Q(xx,a') =')','')='' Ek
Similarly to T, the one-step operators T and T lead to space complexity issue by producing distributions with a number X' of atoms, which can be too demanding for large state space.
3.2 Main properties
Let us now discuss some key properties satisfied by our new one-step distributional operators. We first show that, similarly to the non-distributional sctting, our approach comes with contractive operators in both evaluation and control.
Proposition 3.1 (CONTRACTIVITT). Let be policy.
(1) For any 1 < S co, the one-step operators J" and T are y-contractions in W (ii) The Cramér-projectee one-step operators He o I and Te J are 7-contractions in W.
We stress that Proposition 3.1 highly contrasts with classic DistrRL where the control operator T is not contraction in any metric. A major consequence of contractivity is the existence and uniqueness of fixed points, whose explicil formulas are gathered in the next proposition and in Table
Proposition 3.2 (FIXED POINTS). Let be a policy and consider He from Definition 2.3

---

Published in Transactions on Machine Learning Research (08/2023)
0.8
0.4
0.4
0.4
0.2
0.2
0.1
-04
-0.2
0.0
0.2
0.4
0.8
1.0
1.2
1.4
1.5
1.8
1.00
1.25
1.50
1.75
2.00
2.25
255
275
(a) Po
(211022)
(b) [A2
(r1.92)
(c) P4
(211102)
Figure Illustration of Examples 2.2 3.3 in the two-states two-actions MDP depicted in Figure with discount factor 7 = 2, for the stochastic policy π <alrr) = 1 for all x, a. The probability distribution H4 (0102) with histogram represented in red (resp. in blue) is obtained by applying successive times the one-step distributional operator T" (resp. the "full" distributional operator TT) to the initial distributions (1raa 80- The red "one-step" histogram has constant number of atoms, corresponding to the number of states. Meanwhile, the blue "full" histogram approximates a continuous probability density with an increasing number of atoms along iterations.
(1) The unique fixed point of T is vx given by
vlaa)) P(d' (r)ra
(it) The unique fixed point of T is v. given by
(a.a)
Σ P(ol (raaa)aa
(il)) If S r(z,,,, x') + +zV=(x) S zk for all triplets (x, a, x'), then the unique fixed point of 1lc o J is
77 = [[(()).
(iv) If < r( x') (V) s zg for all triplets (I, a, a.), then the unique faced point of Ie T is
7/- = [.VV...
The proof of Proposition Ex. P(="|=', (r)((((', a', x)) (V*("") with the simpler Q-value Q" (a', a'). This substitution is per- missible due to the unique fixed-poin characterization of V and Q Interestingly, Proposition (.2) (iii)- shows that, in our one-step framework, the fixed point of a projected operator is simply the projection of the fixed point of the unprojected operator. Although this fact seems natural, it is not necessarily true in DistrRL. Indeed, the proof of Proposition is a worse approximation of μą than is lcM by a multiplicative factor VI/(1 - 7) terms of Cramér dis- tance. Furthermore, we stress that two policies sharing the same value function but with potentially distinct distributions and risk-levels (e-g. different variances) cannot be distinguished via the one-step approach as, fixed point. This observation highlights notable limitation of one-step DistrRL: it is unable to distinguish between two policies with identical value functions that would, however, yield distinct fixed points in the context of full DistrRL. Equipped with our projected one-step operators, we derive in the next section variants of the CDRL algorithms.
3.2 (point for instance) involves replacing, inside each Dirac, the quantity
(Rowland et al. 2018) suggests that the fixed point of Ie T
from Proposition 3.2 (i), they necessarily share the
4 One-step DistrRL algorithms
This section introduces new categorical algorithms based on the Cramér projection together with formal convergence guarantees.

---

Published in Transactions on Machine Learning Research (08/2023)
DistrRL
One-stop DistrRL
Evaluation
Distr (2 = 7'r (( e, A,, Xt+1))
P(a le, (er,aaa ()
Control
does not necessarily exist
x' P(x ( (r (1) ((
Table 1: Comparison of the fixed points of the distributional Bellman operators in DistrRL versus one-step DistrRL.
CDRL
Ono-step CDRL
Categorical target
He Σ k= (r,4(1 46)
He (oets(((m))a"))) Olog K)
Time complexity
O(K log K)
Comparison of the categorical targets in CDRL versus one-step CDRL. The notation a" refers to a greedy action, namely a" E arg max Qr(Fe+1,a'..
Table 2:
4.1 One-Step CDRL
We propose two algorithms, for policy evaluation and control respectively, that are described in Algorithm These new categorical methods are derived from the stochastic approximation of the projected operators Ie TTT and Ie J. For each state-action pair (z, a), both methods learn discrete probability distribution pi (1 ....a)....P............. PK (z, a) over some fixed support 21 <.. < zк. As in classic RL algorithms,
consider a sequence of stepsizes a1(x, a) 2 O indexed by states, actions and time steps t 2 0. At each time t, we perform a mixture update between the current distribution and a distributional target which is the Cramér projection of a single Dirac mass located at the target of TD(0) or Q-learning, computed from a single transition (ze, Qe. Pt, 11) The only difference with tabular CDRL lies in the target as shown in Table Contrary the original CDRL target, ours remains the same whatever the greedy action we choose inside the set arg m (1 a'): this explains why our approach is stable even when the optimal policy is not unique as illustraaed in Figure 2 Moreover, our categorical target is faster to compute than in CDRL, where the time complexity O(K log(K)) pays the price for inserting every atom Tt + 7 into the sorted array Zk). Before moving forward to the convergence analysis of Algorithm
we propose its deep
Algorithm Tabular one-step categorical DistrRL.
(z,a) = Ek= Pt, (x, a)8z, for all (z,a) Input: T/t T/
Sample transition: (It, ae, Te,1))))
Estimate Q-valucs: (x(1++1, + Ek Pe, Pe, (2)411))))) 24
if policy evaluation then
(().a)) +
(e = Pe
else if control then
)(.a))) + -ee(d max, (x(me44) = =EK Pr.
end if
Mixture update: m+t1 + (1- (1-a a)) (Z( + arl2, (f
(m),)))
(a.a) (z,a) a(z,a) # (c)
Mh4 +m 1tt
Output: M
RL counterpart in Algorithm
2 based on the minimization of Kulbeack-Leibler (KL) loss.
Deep onc-step CDRL. The main challenge in a non-tabular context is to learn the distributions in a compact and efficient way. For that purpose, we use the same deep categorical approach as in C51 (Belle- mare et al. 2017). As shown in Table [ the one-step method aims at learning much simpler distributions (nocessarily atomic) than full DistrRL. (typically, continuous distributions carrying more information). This suggests choosing a smaller number of categories, e.g. K used in Section 5 compared to what is commonly used in CDRL. Due to its similarity with C51, we call this new algorithm "OS-C51".
8

---

Published in Transactions on Machine Learning Research (08/2023)
Algorithm 2 OS-C51 (single update)
Input:
categorical distributions No Compute Q-function in next state: Q(1441, d') t E P8, (#c+1,
(2xa) Σk Po,k (z, a)ss, and a transition (0t, ae Pry Ft41)
Compute categorical target: sy(er1n)) t- [ max, Q(#241,a'))
Output: KL( mm)))))))
4.2 Convergence analysis
We now provide convergence guarantees for our tabular one-step DistrRL. algorithms. The major difference with the analysis CDRL (Rowland et al.) 2018) that we do not require the uniqueness of the optimal policy in the case of control. We also rely on the existing analysis of non-distributional RL (Dayan 1992; Tsitsiklis 1994 Jaakkola et al., 1993). In particular, require the following standard assumption.
Assumption 4.1. For any pair (r,a) E X' x A, the stepsizes (an(2, a))o20 satisfy the Robbins-Monro conditions:
a) oo 01 a)² <
almost surely.
Equipped with stepsizes (2, a)) satisfying the Robbins-Monny conditions, we are now ready to state our main thcoretical contribution: namely, the convergence of Algorithm
Theorem 4.1. Consider Algorithm (0,(x,a) a)),
i and let us assume that Assumption holds for the stepsizes
(1) In the case of evaluation of a policy
W (mhs Un) 10000
almost surely,
where ny defined in Proposition 3.2 (ii8).
(it) In the case of control,
W (me. n.)
t 0 almost surely,
where n7 is defined in Proposition 3.2 (iv).
The proof of Theorem 4.1 is deferred to the Supplementary Material: it follows the same steps the proofs of Theorem 2 in (Tsitsiklis 1994] and Theorem 1 in (Rowland et al.) 2018). Notably, the same for evaluation and control contrary to (Rowland et al. 2018) where the control case requires additional assumptions (namely, uniqueness of the optimal policy and for all t 2 0, for all 1 S k < K such that ( a²) f 0, rz 7 E [x] almost surely).
analysis remains
5 Numerical experiments
In this section, present numerical experiments on both tabular and Atari games environments.
5.1 Tabular setting
We describe our tabular experiments: first in a dynamic programming context i.c. knowing the transition kernel P and the reward function r, then in the "Frozen Lake" environment by only observing empirical transitions.
9

---

Published in Transactions on Machine Learning Research (08/2023)
10-
1.0-
3000
||Q 10a--'
0.
0.8-
0.6
Palks, Sg)
0.6-
D-0-n m-n---
2000
0.4
0.4-
1000
500
0-
0
20000 40000
50000 80000 100000
0
20000
40000
50000
80000
100000
20000
40000
50000 80000 100000
(a) (Za, 0g)
(b) (T11,
(c) -
Figure 4: Learning dynamics of the atomic probability distribution (1r, d)) produced by Al- gorithm
1 (control) on the Frozen Lake environment across 100000 iterations, with support (21, ø2, 23) =
(0, 10, 20), constant stepsize a 0.6 and e-greody exploration. The results are averagod over 100 scods.
Distributional dynamic programming. Figures 3 are obtained by exact dynamic programming in the MDP in Figure 1 with 7 3: in this specific case, all policies are optimal. As discussed in subsection 2.3 this handerafted example reveals the instability of classic CDRL in Figure 2 contrary to our one-step approach: for both methods we consider K = 4 categories with (21, (2,,22, 23. 21) = 1.9, 10). We point out that our results are sensitive to this specific choice of support. Indeed, always taking the action aj in state 21 results in a total discounted distributional return reduced to a Dirac mass at 2. However, taking 2 but with non-zero variance. Thus, CDRL
the action 02 leads to the same expected value Q" (21, targets can occasionally (when taking action 02) fall outside the selected interval [1.9, 2.1], leading to the depicted instability. Figure 3 illustrates the space complexity issuc of DistrRL without projection: while the number of atoms is multiplied by at most | after each application of T", our one-step operators produce distributions with (at most) as many atoms as there are states, i.e. two in this example.
Frozen Lake, We also consider the Frozen Lake environment from OpenAI Gym (Brockman et al. 2016
with discount factor 7 = 0.95. It is characterized by 16 states 21,...,,31 T16 and four actions d1,..., Q4. Plus, this is a stochastic environment i.e. the transition probabilities P(œ'Ix, a) are not all equal to either 0 or 1, which justifies a distributional approach. The FrozenLake environment is a gridworld game (4x4 grid), where an agent. moves on slippery ice (stochastie transitions) towards goal. The agent can take 4 actions (up, down, left, right), aiming to reach the goal state from the starting state, avoiding holes, and only achieving reward of 1 upon reaching the goal, and 0 otherwise. Due to the slippery attribute of the environment, the transitions are stochastic, meaning the agent may not always end up in the state it intended to move to. For example, if the agent chooses to move right, it may actually move up, right, or down with equal probability 1/3 due to the slippery ice. In Figure ! we plot the iterates Pt, a) over 100000 steps generated by our tabular Algorithm [ with K = atoms and 22, 22, 23) = (0, 10, 20), constant stepsize a = 0.6 and e-greedy exploration with E exponentially decaying from 1 0.25. We average the results over 100 seeds. Given a pair (I., a), observe the joint convergence of the three probabilities. As in CDRL, and because of the mean- peen-ppeserving property (Eq. 2), the average Q((2,a) Σi= Pt, a)æk coincides with the Q-learning iterates and converges to Q*.
5.2 Atari games
For the experiments on Atari games (Bellemare et al. 2013), we implement the OS-C51 agent on top of the popular "CleanRL" codebase (Huang et al.) 2022). We compare the OS-C51 algorithm against C51 for two different number of atoms: K = 4 or K = 51. In each of the two cases, we choose the support to be evenly spread over the interval [-10, 10]: 110 < 2ボ 10 We use the same architecture C51: deep neural network, parameterized by 6, takes an observation as input and outputs a vector of K logits. We run DQN and C51 with the default hyperparameters from CleanRL; for our new OC-C51 method, we take the same hyperparameters as in C51 due to their similarities. All three methods are based on a deep neural network with 3 convolutional layers followed by fully connected layers with ReLU activation functions.
4 herr:: https: www.ichhbb.ccm/mstttaa//cceaaaaa
10

---

Published in Transactions on Machine Learning Research (08/2023)
BeamRiderNoFramesxipvv
BreakoutNooFrmmesxipvvv
PongNoFrameskip-v4
500
10000
8000
Episode score
400
Episode score
300
6000
Episode
200
4000
DQN atoms]
DON atomm) cs1 (51 atomms
DON
-10
(4 wtoms)
2000
Cs1 [51 wttoms)
C51 (51 stomsl
05 C31 (4 stoms)
C51 14 stoma!
05 C51 (4 atoms)
05-t5i (51 atomel
05-C51 151 atoms)
-20
05 351-511 (S1 atoms)
0.0
2.5
5.0
7.5
10.0
0.0
2.5
5.0
7.5
10.0
0.0
2.5
5.0
7.5
10.0
Million frames
Million frames
Million frames
(a) Beamrider
(b) Breakout
(c) Pong
Figure 5: Average evaluation episodic return over 10 million frames for three Atari games.
They all use the same linearly decaying e-greedy exploration and the transitions are collected in a replay memory buffer of size 1000000. In particular for DQN, we use the Adam optimizer (Kingma & Ba 2014) with learning rate equal to 0.0001 and batch size of 32. For the optimization of C51 and OS-C51, we use the Adam optimizer with learning rate set 0.00025 and batch size equal to 32. The results are averaged over 5 seeds. As shown in Figure [ the performance of OS-C51 (with only four atoms) is comparable with C51 on the Beamrider and Pong games. However, there seems to be significant advantage for C51 on Breakout: this could be attributed to C51's objective of approximating distributions with more complex and richer information, potentially offering a deeper understanding and better representation of the game dynamics. Finally, we highlight that we did not use "sticky actions", where a chosen action is randomly repeated for several consccutive frames, introducing a form of environmental stochasticity. This could be an interesting source of stochasticity to explore for future experimental investigation, potentially adding complexity and richness to the learning process.
6 Conclusion
We proposed new distributional RL algorithms that naturally extend TD(0) and Q-lcarning in the tabular setting. The main novelty our approach is the use of new one-step distributional operators circumventing the instability issues of DistrRI. control. We provided both theoretical convergence analysis and empirical proof-of-concept. A significant limitation of one-step DistrRL, as underscored in our study, its inability to distinguish between two policies with identical value functions but differing distributions. This points to a crucial advantage of full DistrRL, which can capture and exploit these subtleties. Future research could investigate the generalization of our method based on the dynamics of several succossive steps instead of single one. Investigating this multi-slep extension would likely reveal further important insights inlo the dynamics of RL processes. Another promising direction for future research lies in exploring the possible connection between full DistrRL and a multi-step approach with infinitoly many steps. Such an investigation could offer deeper understanding of the relationship between these methods and help to clarify the distinct benefits and trade-offs involved.
References
Mastane Achab. Ronking and risk-aware reinforcement learning. PhD thesis, Institut polytechnique de Paris, 2020.
Mastane Achab and Gergely Neu. Robustness and risk management via distributional dynamic programming. arXiv preprint Xrx::21122 15430, 2021.
Mastane Achab, Reda Alami, Yasser Abdelaziz Dahou Djilali, Kirill Fodyanin, Eric Moulines, and Maxim Panov. Distributional deep q-learning with ovar regression. In Deep Reinforcement Learning Workshop NewrIPS 2022, 2022.
11

---

Published in Transactions on Machine Learning Research (08/2023)
Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253 2013.
Marc Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforcement learning. In International Conference on Machine Learning, pp. 449 458. PMLR, 2017.
Marc G Bellemare, Nicolas Le Roux, Pablo Samuel Castro, and Subhodeep Moitra. Distributional reinforce- ment learning with linear function approximation. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2203 2211. PMLR, 2019.
Marc G Bellemare, Will Dabney, and Mark Rowland. Distributional reinforcement learning. MIT Press, 2023.
Richard Bellman. Dynamic programmingg Science, 153(3731):44 37, 1966.
François Bolley. Separability and completeness for the wassenstein distance. In Séminuire de probubililés XLI, pp. 371-377. Springer, 2008.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiu:1606.o 0100000..44 2016.
Will Dabney, Mark Rowland, Marc Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018.
Peter Dayan. The convergence of td (A) for general X. Machine learning, 8(3):341 362, 1992.
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mchta, and João G.M. Araújo. Cleanrl: High-quality single-file implementations deep reinforcement learning algorithms. Journal of Machine Learning Research, 23(27)): 2022.
Tommi Jaakkola, Michael Jordan, and Satinder Singh. Convergence of stochastic iterative dynamic pro- gramming algorithms. Advances in neural information processing systems, 6, 1993.
Diederík P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprinl arXiv:1412.6900 2014.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglon, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXie preprint urXiv:1312.5602, 2013.
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):559--5)) 533, 2015.
Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka. Non- parametric return distribution approximation for reinforcement learning, In ICML, 2010a.
Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka. Para- metric return density estimation for reinforcement learning. In Proceedings of the Twenty- Sixth Conference on Uncertainty in Artificial Intelligence, pp. 368-375 2010b.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Mark Rowland, Marc Bellemare, Will Dabney, Rémi Munos, and Yee Whye Teh. An analysis of categorical distributional reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pp. 29-37. PMLR, 2018.
Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Haru- tyunyan, Karl Tuyls, Mare G Bellemare, and Will Dabney. An analysis of quantile temporal-differenn learning. Xixio preprint arXiv:331 arXi:2301.004622 2023.
12

---

Published in Transactions on Machine Learning Research (08/2023)
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3(1):9 44, 1988.
Richard S Sutton and Andrew G Barlo. Reinforcement leurning: An introduction. MIT press, 2018.
John N Tsitsiklis, Asynchronous stochastic approximation and q-learning, Machine learning, 16(3):185 202, 1994.
Christopher JCH Watkins and Peter Dayan. Q-laarning. Machine learning, 8(3):279 292, 1992.
A Proof of Proposition 3.1
i). Contraction in Wasserstein distance. Tet 1 s Let 11, E Po(R) be two distribution
functions. Denoting Qi (2, a) and Q2(2, the respective expectations of μl (1.a) and 11 (x.0), we have for any pair (z,a):
WP((TTI)) (1.a), (Jnz)(a.a))=
WPC
P(x (c, (0 max, Qi(dla)); E P( la,, (),,)))0r0 max. Qa(a' aa))
1
P(a' az a) a)W
max, ( (8 max, Q,lrl,") or(z, max,, Q2(=',a1))
7p Σ P(a' x a) max Q1 d') max (2(x',a'))"
Σ
P(x (ellzoo (x)a) a) max 1Q2 a') - 12(2222'2222 = 70 Σ
max
- Fa P
1=0
S <> Σ P(xlIz,, max
FT (T)
Frld S " (M1,
(y als 11 (0))
where Fu,o; Fa.o denote the respective CDFs of μi interpretation of the Wasscrstein distance as an infimum over couplings. Indeed, we have upper bounded
(a,a) and pe (l)a)) and the first inequality follows from the
(F((T))) (Tuz)(ral)) inf (Zz,,zm>all 7
AEA
where A denotes the set of all couplings of (T,)) and (T via specific coupling, Ao such that
>o((r(z,a,z)) x') + max Q (z', x (r(x, x') + max (2(2' P(al(a, a), Ve'.
Hence, by taking the supremum over all (x, a), we deduce that T is a y-contraction in W
W (J Tu2) S W WW (11,
13

---

Published in Transactions on Machine Learning Research (08/2023)
Similarly for p oc,
Wo ((TN)) ((r,)) (Jma) (Juz) (x,))
Wx
P(xl (x, ())r((aaz)))))))))))))) 01(2 E P(x' (00 Qz(x',a')
P( [x, a) W 2o (o max, Qi(e' (aae),rr((eee ar"a max Qa(a',a'))
r Σ P(a' x a)) max Qi(x'aa)) max 22(x',a))l
a²
f (F -a (T))
$ Σ Pla' p. a) max QQ 12('', =- > P(a' (1) P(
=
max
1=
S y Σ (( max [ F
(T) (a)
F- (T) S o (11)
For policy ovaluation, one can follow the same proofs by incorporating the following additional stop:
Σ (dd1 s Σ x(d''=') Q2(2', s max Qi a') Q2(2'2 (2))
a'
ii.. Let us write:
(1eJJ ((((UCTR)) (F
(m)))
W1 Σ P( Ir, a
max Qi(ataa)) Σ P(a' I
(a(a%)a)
$ P(a'lc, a) W/ ( max w'), 1 max, Qz(e',a") a())
5 Σ P(a' x a) max Q1 max Q2(x',a))
a²
Σ Plx' [z, a) max 10)(1', a') 20)'''a'11=1 = Σ P(a'lz (( a) max /' (F)
<>
a'
1=0
{ F (T)
S Σ (( max
(T) S wW (21,)1))))
o'
Wi Wi(H (a))
where the second inequality follows from Lemma proof. The proof for HT is similar.
A.1 Taking the supremum over all (x, a) concludes the
Lemma A.1. Let a, b be two real numbers. Then,
Wi WT Mcd.) S |a b).
Proof. We proceed by exhaustion of all possible (redundant) cases, up to permuting a and
(i) a, b are both outside the interval =x]= the proof is trivial in this case
(ii) e (x1, b E (xk and < b
(iii) a E (k,, b E (,, and a > b
(iv) (21, and & (21,
14

---

Published in Transactions on Machine Learning Research (08/2023)
(v) a, b are both inside the interval [x,
Proof in case ii). We assume that b belongs to the support, b 2 a and a lies inside the interval (21,
< a S 2j41 and b b= 2k+1 with i K.
Then, we have:
W1 WIllloo 11b6,)=
3j+1 (ZR+1 2z) + 2j 2j+1 2j
(2x+1 2z+1)
2j+1 - 2j
(#x+1 - 21) 2/+1
2j
(2/+1 - -2/) a a=
2j+1 2j
2j+1 Z2
= =k+1 - 2j (a 2z} = -k+11 == aa b/.
Proof in case iii). Similar to (ii) except that we have b b= = 2k with j 2 k:
Wl 11c66)=
2j+1 2j (2g 2/:) +
3j+1
2j
(2j+1 2x))
2j+1 2j
2j+1 a + - 2j + (2j+1 - -2z) a-3 2g
= (z) 2kk)
Zj+1 - Zg
2j+1 2j
=2 2x + a 2j = a 24 = a b/.
Proof in case iv). Now, if is outside the interval (21, say S 21, we deduce from the previous computation that W1 ( Ieo,)) = a 21 $ la b). Proof in case v). Finally, if both and b lie in (21,
(Ileo,,, m16+)) Wi( 3-1 a
2/+1 2j 6z,- + a 8 8eff1: 1l(6,)
2j+1 2j
< 2j+1 Wi(0,, I1) + a- 2,
W1 W6 Ilc
2j+1 - Zj
24+1 - 2j
2j+1 - a b - 2z] +
a - 2g b b- 2541/1 (3)
2j+1 Zg
Zj+1 2j
where we used cases (ii) and (iii) in the last equality. Then, if 2 1j+1, Eq. equation implies
3j+1 a (b - 2n) +
- 2j (b 10-a0 2j+1)
WIlCo,,, (1cda) S
2j+1 -
Zj+1
(b 2j) Zj+1 - a + a -2j (2j+1 - -2ÿ) - Zj = b 2j a + 2j b
2j+1 2j
2j+1 2j
Symmetrically, if < 2j, Eq. equation 3
implics
Wl Tedp)
< 2j+1
(z, - b) 2j 21-1
(zz+1 - b) =
2/-1 2j
Lastly, if a. b both belong to the same segment (zj, 11 and say S b, we have by direct computation:
Wi Wl (cd6) = (zj+1 22y) 2++1 - a - (zj+1 - b) 6 --6 a
žÿ Rj
15

---

Published in Transactions on Machine Learning Research (08/2023)
B Proof of Proposition 3.2
Given the completeness of the Wasserstein space (see Bolley 2008 , by combining Proposition 3.1 with
Banach's fixed point theorem, we deduce the existence and uniqueness of the fixed points.
For (i); we verify:
(Tuz) (x,a) Σ (() E,, x(a'(x') E. ""
Σ p( (00,0)(00000'))))))))))) "1 (001))00000) (/
ii). For the control case,
(Tv,) (a,0) Σ (a,aa))))) max Ene P(a"Ia "o"" a' ,0 --""")
= Σ P(el (a Q'(c',a") u (2.a)
iii.. We start by computing:
(T o [lCks) (r,a)
(x,a)
Σ P(x' a)o
(2, E,, ((d111)0002" vla,a)
where used the mean-preseroing property of the Cramér projection (see Eq. equation Then, deduce that
Il a JH Il IIe 7
iv). Similarly, one can show that n, is the fixed point Ile o T.
C Proof of Theorem 4.1
We follow the same steps in the proofs of Theorem in (sitsiklis 1994] and Theorem in (Rowland) jtt al. 2018). We focus on the control case (ii), though the proof remains similar for policy evaluation. Tet us define for each (xr a) E X' A: L.0 (2.a) 62, Ufz.a) = Egg, and for j 2 0,
(2,a) 1
(1,0) + (1LJL,) (mna) and Ul(e) =
1
1 y + (L,T
We now state the two following lemmas and then, before proving them, explain how they allow us to conclude.
Lemma C.1. (i) In terms of entrywise stochastic dominance, holds for all 2 0: Lit1 2 Lg and
Us+1 S U-- (ii) (Lz) and (Uz) both converge to Tr in WT.
Lemma C.2. Given j 2 0, there exists a random time Tý 2 0 such that
L S n SU for all t 2
T3 almost surely.
From Lemma C.1 (ii), let > 0 and take j j2 large enough such that max m.). W (U,,,, <c. Then, it follows from Lemma followed by a triangular inequality that W (m,, T.) S WW(TT, U,) W WI(T,,n)) + W WU,,, 26,
for all t 2 T,. almost surely. Finally,
W WiGns q-) S Wi(ne L,) W (Ly,
W WW((U,)))) < 5E,
which implies Theorem 4.1 We still need to prove Lemmas C.1 C.2
16

---

Published in Transactions on Machine Learning Research (08/2023)
C.1 Proof of Lemma C.1
(i). Let show by induction that Uf+1 S Uy. First, the base case Uliaa) S U@2,0) ò=x is true because U((,,) is supported in f21...., [k) Now, let us assume that Us+1 I Ug for some 2 0. We recall from Proposition 5 in (Rowland et al. 2018) that Ilç a monotone map for element-wise stochastic dominance. Plus, it easy to see that both J" and T are monotone too. Then, by monotonicity of IL,T it holds that HUJU,+1 I IC and hence,
Uff2
I fufoaa) + ( (T j+1 2 (ufer) 1 (TTCTUH+)))))) &
which proves the induction step. Symmetrically, one can show that Ly+1 2 Lg
(ii). Similarly to Lemma 7 in (Rowland et al. 2018 formulate the following general result.
Lemma C.3. Let (we)k= be a sequence of probability distributions over {z1,..., 2x] such that vk+1 S V/ for all k 2 Then, there erists a limit probability distribution Vim over { such that vk + Vlim in W1.
Proof. Let us denole Fs the CDF of us and FF its quantile function valued in (.1,..., The assumption that P& stochastically dominates Vk+1 reformulates as (T) S FF (r) for all < S Hence for each = E (0, 1], the sequence FE" (T) is non-increasing and lower bounded by #1. Therefore, this sequence converges: Fx (T) + (T). As Fr (T) can only take discrete values, there exists NN()) such that vk 2 N(T), F (T)
(T) The limit function Him thus non-decreasing, valued in f.i,..., and right-continuous with left. limits. It is therefore the quantile function of probability distribution Plim supported over {=1, (21.-..,,2)) =x Let now show that Wi (Vk, Mim) O. Denoting Fimm the CDF of Vin and po = 0, Fim (21), ... PK = Fim (2x) = I, we have:
(T)
Wi Mim) [" 1 H1m (T)(d)
J T
Σ
IFT - 22/dT.
(4)
ISkSk:pu
Let Ap = minggg (pk pk pk-1), A, 2K - 21 and < 4p/2. Then for all 2 maxiçk< (Kptskk max N (P N ((p - c)}, (T) is constant equal to 2 for any T E [pk-1 + c. c) and we have:
Wi(kk, Plim)
Σ
(T) 2x/dT
IeKSK:pePPe
T=PA-1
itc
Σ
FF (T)
zk dr +
(T) zu /dt
Fr (T)
2k dr
T= i
2A.
<d.
S 2K Aze.
which proves the result.
By applying Lemma C.3 to the sequence (U) )220 for each pair (z, deduce the convergence of (Ux) towards some limit distribution function Unm in W, Finally, by continuity of IeT for the metric W1, this limit must verify:
1
Unm
- Unm + - TU
Ulm
2
from which deduce that Ulim no by unicity. The proof that (Igz) converges to n,. in W is analogous.
17

---

Published in Transactions on Machine Learning Research (08/2023)
C.2 Proof of Lemma C.2
Let prove Lemma C.2 by induction. The base case j 0 is true because for all t To 0, m/ supported on (.1,..., 2 and so:
(ara)
is
T.0 (x.0) = =ō. S (2,0) S 62N = U(a.))
Then, let us assume that the result true for some j 2 0, i.c. there exists random time Tj such that Li in U; for all 2 T almost surely. Now, let us show that there exists random lime Tj+1 such that Th S Us+1 for all t 2 Tj+1 almost surely ( the proof for n 2 Ls+1 is analogous). For each state-action pair (xr, a), we define Hy, = Ular)) and Wiz the zero measure i.e. Wi, (A) = 0 for all Borel sets A C l.
Then, for t T
H( = (1 -1-ar(x,, a)) H)) (1)aa + (z (T and
1 = a)) + anla (Te (f(2,.)) max. Q=(x",a))) (Tc.))))(aa)))))
(5)
where I' ~ a)) Q0(2', = Exn pLo(e', (((r',d)) d')zx. Moreover, (x, = (xz, then = xt-1 and r( ae x241) = rr. A consequence of Eq. 5 is that w(aa)) is signed measure R with total measure equal to zero: (a, (R) = 0 for all 2 Tz. Let us now prove by (another) induction that n (2,a) < (x(0)) + w(
for all 2 Tj, (z, a) E x A almost surely. For = Tj, it is true by assumption that Tr, S Uy almost surely. Then, suppose that n z He + W, almost surely for some t 2 Tj. Recalling that ax a) for all (z, a) & (x,,a)) a)), it holds that:
(x.a) (1 al(x,a)) (=0)) + all(,,)lle max, (c(r((a))) Mbt1
(1 az az)m (2.a) a (2,) + az(x, a) (6 (0,(2, max, q,(x',aa) (ore)) WT) (ao) ( (2,)) + a1 a) (e (6 (2,a, max, (=(w'aa') (TeJin) (mse))
< (1 (1-ou
()) (11-
= =H(aa) wlaaa)
(6)
where the inequality comes from the assumptions S +1++ W. and ne S Uj combined with the monotonicity of H. Now, obscrve that He can be explicitly expressed as follows:
II (1 or (1, U, + (1 (1- II (1 (1-ar(z,))) (I.JU))(a.)) T==T,
((z.a)
U==;;
Since by Assumption 4.1 it holds that Ee a2((2, 0o for all (x, almost surcly, deduce that there exists a random time T)+1 2 Tj such that I (1 (1- av- a)) S for all (z, a) and for all 2 T/+1 almost
surely. Then, because Lemma C.1 (i) implies that NeTU I Uj, we have for all t 2 Tz+1:
1 U, + 1
3
LCJU + w, 1 U, - JU, + W, 1 (U, -IIJU,)
Th < H, + W. <
4
2
-U3+1 W. 1 } IcTU,). (7)
We point out that the random noise term appearing in the definition of w(fx,)) has zero mean: for all x+1
1 k k
ExwP()a)) llo [o((rr(..az')t max,, Q((r(aa)) (N.TTn (m.a))
((- ((-00, (x)) =0
(8)
Eq.
8 implics via a classic stochastic approximation argument under Assumption
4.1 that
((-00, -> almost surely, for all (x, a) and k E {1...., K). Finally, take Tjt1 Tz+1
18

---

Published in Transactions on Machine Learning Research (08/2023)
large enough that a1]1 S _ for all & 2 T3+1 and all (x, a), almost surely, where A is given by:
= inf | - ( 211)
# fo:
Indeed, if Uf (1 zx]) (IcJU, (( ((,00, 2k]), then
which concludes the proof by using Eq. 07 ((-00, 2x]) U3+1" ((-o0, 2x)).
19


Reviews:
[{"review_id": "W0d9nqUaX1", "reviewer": "Reviewer_SmEz", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\nThe paper is very well written. I particularly liked section 2 and the discussion in section 4.1.\nI also appreciate the motivation for this work. Even though it's a simplification of the original DistRL algorithm, this enables theoretical results that were previously not possible (although see questions below).\nThe empirical evaluations are also effective at transmitting the message the authors are trying to convey.\n\n# Weaknesses\n\n## Source of value functions\nMy main concern with this paper is around Proposition 3.2, which has implications for Theorem 4.1, the main theoretical result of the paper.\nIn the statement of proposition 3.2 the authors define $\\nu_{\\pi}^{(x, a)}$ and $\\nu_{\\*}^{(x, a)}$ with respect to $V^{\\pi}(x')$ and $V^*(x')$, respectively (also mentioned in Table 1). Where are $V^{\\pi}$ and $V^*$ coming from? Are they the value functions you'd obtain with (standard) expectational RL? If so, this appears to be problematic, as they would _not_ be the same values you'd get with the proposed one-step method. Indeed, the proof of Proposition 3.2 seems to rely on $\\nu$ being defined with respect to these value functions.\n\nGoing through the proof, it seems like the value functions can basically be replaced with _any_ function $f:\\mathcal{X}\\rightarrow\\mathbb{R}$, as long as $\\nu$ is defined with respect to $f$ in the proposition statement. In other words, there is no need for the true value functions for the proof to hold.\n\nAnd this leads me to question the significance of the results. It seems they're really mostly concerned with immediate reward, but are basically agnostic to downstream implications that are affected by the environment (see comment above about replacing with any state function). This makes sense, since it is one-step after all, but then I'm not sure on the claimed strength of connection to the original distributional RL results.\n\nThe empirical results suggest otherwise, but I'm wondering if this is a consequence of the environments rather than the proposed algorithm. A baseline that would be more convincing is to run a one-step DQN, say. This would verify if dropping the Bellman backup still results in good performance. If so, then I would search for environments where this is _not_ the case. Also consider adding sticky actions (if you're not using them already).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents a variant of distributional reinforcement learning (DistRL), where the standard Bellman backup is replaced by a one-step backup. The advantage of this approach is it allows proving convergence in both the policy evaluation and control cases, without requiring assumptions such as uniqueness of the optimal policy (as required in prior work). Through empirical evaluations the authors also present evidence that their simpler method can reduce training instabilities in certain edge cases, and has comparable performance in some standard benchmark environments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The main point to be addressed are:\n1. The one pointed out above, with regards to the use of value functions in the theoretical analyses. This is a critical point to address.\n1. The conclusion is a bit terse. I think it could be expanded a bit more, in particular with respect to limitations, practical implications, and possibilities for future work.\n1. In the proof of Lemma A.1 the second inequality for case ii doesn't seem to follow. The first line expands to\n$z_{j+1}z_{k+1} - z_{j+1}z_j - az_{k+1} + az_j + az_{k+1} - az_{j+1} - z_j z_{k+1} + z_j z_{j+1}$\n\nwhereas the second line ends up with an extra $-z_j^2$ term.\n\nSome other minor points:\n1. Figure 1 could be moved further down, it's a little distracting where it is, and not needed there.\n1. You should cite the recent [Distributional RL book](https://www.distributional-rl.org/), by Bellemare, Dabney, and Rowland. In particular, there's a counterexample in chapter 7 showing that one can't remove the condition of unique optimal policy to guarantee convergence in the regular DistRL setting.\n1. In Figure 3:\n\n    a. Add a legend\n\n    b. Add the discussion on growing number of atoms in the caption\n\n    c. Remind the user the difference between $\\mathcal{T}$ and $\\mathscr{T}$\n1. In Algorithm 1, where are the $p_{t,k}$ updated? The $t$ subscript suggests they should also be updated in the loop, but it is not specified.\n1. In Figure 4 are you only showing OS-CDRL? How does CDRL compare?\n1. Add more context to the caption of figure 4; it's best if figures are (mostly) self-contained.\n1. Why were $\\gamma=1/2$ and specific categories chosen for the DP example in 5.1? Are the results sensitive to this?\n1. In the DP example in 5.1, could one argue that the increased number of atoms results in more expressivity of the distribution, which is necessary for regular DistRL, but not necessary for one-step? Put another way, is it just that full DistRL results fin a richer distribution that requires more atoms?\n1. In the Frozen lake discussion you seem to be suggesting a distributional approach is not justified in deterministic settings. Is this necessarily the case?\n1. In 5.2 are the same hyperparameters used for C51?\n1. In 5.2 are you using sticky actions?\n1. Please restate proposition/theorem statements in appendix", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "aUGGDkhfjb", "reviewer": "Reviewer_Q3pB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "positive points:\n\nThe overall writing is clear.\n\nThe technical results look solid to me.\n\nThe claim of contribution is clear and the results support the claim well.\n\nnegative points:\n\nThe paper does not clearly studies what we lose by considering the one-step distribution instead of the full return distribution. To appreciate the importance of the new algorithms, this point needs to be clearly said.\n\nThe paper does not provide insights into why existing distributional control algorithms fail to converge and how the new algorithm overcomes this issue.\n\nThe empirical results lack experiment details (like what hyperparameters have been tested and how the curves in the figures are generated).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes two new distributional reinforcement learning (RL) algorithms and studies their convergence properties. The two algorithms include a prediction algorithm and a control algorithm. The main contribution is that the new control algorithm can be shown to converge while the existing distributional control algorithms can not. Technically, this contribution is not novel because it mainly uses tools used in previous works. To further verify this convergence result, this paper also shows empirically that the control algorithm indeed converges in a small-scale environment called Frozen Lake environment, thus verifying the convergence result. In addition, the neural network extension of the control algorithm has been tested in three commonly-used atari domains. The empirical results of the new control algorithm seem to be comparable with a state-of-the-art distributional control algorithm.\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "see weaknesses", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xbGHmK3U54", "reviewer": "Reviewer_J5Ja", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength\n---------\n\nThe trick is simple, does not affect the optimality and looks even simpler to implement than classical DistRL.\n\nThe math are solid, and making the projected Bellman operator contractive is a strong improvement.\n\nThey show a very simple example of 2-states MDP where the classical approaches fail while that one converges.\n\nWeaknesses\n--------------\n\nThe experimental part looks a bit sloppy, with only 3 games and results averaged over 3 seeds.\n\nThis makes the result very hard to interpret and I would not say that the new method does not deteriorate the performance of classical DistRL.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a new approach for distributional RL (DistRL) in which the target distribution is modelled by a Dirac of its average (instead of the learned distribution at the next step in classical DistRL).\n\nThat simplification makes the (cramér-projected) Bellmann operator contractive and solves the instability issue in case of non-unique optimal policy. \n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Critical changes\n------------------\n\nThe paper needs more experimental justification. At least averaging over 5 seeds and much more games (the ideal would be all the 57 Atari games). \n\nOther changes\n-----------------\n\nOne important ingredient for the Banach's fixed point theorem is that the metric space is complete, which is never specified.\nIn the case of Wasserstein metrics on probability distributions the completeness is known (for example: Bolley, François. \"Separability and completeness for the Wasserstein distance.\") but I think a sentence should mention that as it is not a well-known or trivial result.\n\nIn def 3.1 (and 3.2 and next) it took me a while to understand that all the target terms are in the Dirac (and not just the reward). I guess just adding parenthesis would help the reader a lot, or even a notation \\delta(x) instead of \\delta_{x}.\n\nIn fig. 4 (Frozen Lac) it would be great to explain a bit what are theses states and what the different actions mean.\n\nIn Appendix proof of prop1, I did not understand the sentence \n\" where the first inequality follows from the interpretation of  the Wassertein distance as an infimum over couplings \"\nIsn't the first inequality a simple Jensen (W_p being convex) ?\n\n", "broader_impact_concerns": "Does not apply", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper identifies a problem in the theoretical application of the most common form of distributional RL. It proposes a new one-step rule, supported by a solid formal analysis (including a convergence guarantee without the assumption of a single optimal policy in both control and evaluation settings, along with a demonstration in a small domain) and several empirical results.

There remain several outstanding points that the authors should address: 
- all minor clarifications remaining from all reviews and responses, 
- the outstanding clarification to Proposition 3.2 pointed out by reviewer SmEz

I also encourage the authors to add one experiment using sticky actions to improve the paper suggest by SmEz, which would improve the empirical evaluation further.

Decision: Accept (poster)