We thank the reviewer for the detailed assessment of our work.

We appreciate your comments about the scalability of our approach and the motivations for studying foundations of multivariate DRL; we have discussed these in detail in the general response.

Q1: The reviewer claims that our work is incremental given the results of [ZCZ+21]. We respectfully disagree, and we hope the following clarify our novelty:

1. [ZCZ+21] only analyzes the contractivity of the unprojected distributional Bellman operator. This result alone does not prove convergence of any practical algorithm. Our work is the first to provide analysis for the type of distributional operator that we can deploy in practice, and associated TD-learning algorithms, unlike [ZCZ+21].
2. [ZCZ+21] only provides convergence analysis for dynamic programming (again, with the unprojected / intractable operator), which we also generally cannot perform in practice without knowledge of the MDP. Our work goes beyond this and provides convergence guarantees for TD-learning, which can be computed from samples of the MDP (and no knowledge of the transition kernel or reward function). This is a much more difficult result to achieve, and we actually require new proof techniques and models to accomplish this (e.g., we must represent distributions as signed measures). Our novel analysis inspired a new algorithm which has not yet been studied even in the $d=1$ case, and is the first provably convergent TD-learning algorithm in the $d>1$ setting.
3. While we have done much more than extend the analysis to MMD, this itself is a bonus. There are fundamental difficulties to performing TD-learning in Wasserstein space due to biased gradients. This is another facet under which our analysis applies to a practically-relevant algorithm, unlike [ZCZ+21].

To be clear, [ZCZ+21] is fantastic and was certainly an inspiration, but our results are not simply incremental modifications of theirs.

Q2: Projections are necessary to avoid exponential blowup in the number of particles in return distributions (see Q4) in DP methods. Some TD algorithms do not explicitly compute projections, just passing gradients through a fixed representation, such as equally weighted particles. However, to prove convergence of such algorithms, one must show that these updates track applications of a projected operator (see e.g. [RMA+23; Sec 5-6]). Crucially, convergence of TD algorithms does not follow from contractivity of the unprojected Bellman operator, and analysis of projected operators is essential. Thus, while [ZCZ+21] does not explicitly compute a projection, this is one sense in which their paper does not theoretically justify their algorithm.

Randomized projections uncommon in the literature: this is a novel contribution that we introduced to tackle issues with computing the EWP projection for DP. In the case of $d>1$, the MMD projection onto the space of EWP representations is non-convex, and could be extremely expensive (if at all possible) to compute in practice. The randomized projection that we introduce allows us to design a tractable algorithm for approximating the multi-return distribution function with EWP representations, which enjoys dimension-free theoretical convergence bounds.

Q3: As described above, our theoretical results are not simple corollaries of existing results. They required novel proof techniques and algorithms (Sec 4-6). Each step of the algorithms described is objectively computationally tractable. With our statements on line 33/67, we are comparing against existing analyses of multivariate distributional RL which rely on regression oracles (for which no tractable algorithm generally exists). See also our general response, which demonstrates that our proposed TD-learning algorithm can be straightforwardly scaled with neural networks.

Q4: Imagine an MDP with $n$ states and a nonzero probability of transition between any 2 states, and $\eta_0(x)$ is supported on 1 point for each $x$. Then, after one iteration of DP, $\eta_1(x)$ will be supported on $n$ points (contributions from each of the $n$ successor states which have return distributions supported on $1$ point). Now, $\eta_2(x)$ will be supported on $n^2$ points, since it mixes contributions of $n$ return distributions each having support on $n$ points. Continuing, we see that $\eta_k(x)$ is generally supported on $n^k$ points.

Q5: We will make this more clear in our revision; the two ideas are modeling return distributions and modeling multivariate returns.

Q6: Good question. Cumulant here refers to the multivariate return. While this is standard terminology in this niche of RL research (see e.g. Bootstrapped Representations for Reinforcement Learning, Le Lan et al. 2023), it is worth clarifying as you point out.

Q7: Alternatively, one could model the positions of $m$ particles and their probability masses, as opposed to modeling empirical distributions. This requires more memory and is often more difficult to optimize (see [BDR23]). Can you clarify in what sense the EWP representation is "inaccurate"? While modeling atom probabilities provides more flexibility, EWP representations are valid probability measures, and we provide strong bounds on the quality of our EWP approximation to the true (nonparametric) multivariate return distributions (Theorem 3).

Q8: We would be happy to explain this further, but it would help us if you could specify more precisely which part is difficult to understand. The proof of Lemma 1 explicitly constructs the QP to be solved for computing the projection. This QP is described completely in the pseudocode on the line with QPSolve -- this is a minimization of a quadratic over a convex set. We can use efficient QP solvers to solve this. It also does not preclude function approximation, since we only need to solve the QP for the target returns (which gradients are not propagated through).

Thank you for your detailed review of our submission. We appreciate the time you've taken to evaluate our work. In our rebuttal, we addressed your concerns regarding the novelty of our theoretical results and provided clarification on how our work differs from existing literature. We also included results from larger-scale experiments as per your suggestion. If you've had a chance to review our rebuttal, we'd be interested to know if our explanations and additional results have helped address your concerns; and if so, we would be grateful if you could consider increasing your score. If you have any further questions, we are eager to discuss them.

We thank the reviewer for their thorough assessment and for their interest in our work. We hope the responses below address the queries raised in the review, please let us know if you have any further questions.

Q1: With regard to $\leq$ vs $\in$ in Theorem 3 -- this was a stylistic choice. We were invoking the definition of $\widetilde{O}(\cdot)$ as a set-valued function -- all elements of this set satisfy the upper bound that you are describing. This is the convention used, for example, in the classic CLRS Introduction to Algorithms text; its meaning coincides with what you are describing.

Q2: You are correct that there is not a unique fixed point of the projected operator. The theorem is simply claiming that, after enough ($K$) iterations, we can guarantee that the resulting $\eta_K$ will be within a very small margin of error from $\eta^\pi$ w.r.t. the MMD. Notably, $\eta^\pi$ is the unique fixed-point of the (unprojected) distributional Bellman operator, which is well-defined (by Theorem 2). Convergence results for dynamic programming algorithms without unique fixed points also exist in previous works, such as [WUS23]. Note we do not require exponentially-many atoms; the free parameter here is $m$, not $K$. Therefore, for a desired tolerance level, we require polylog-many iterations ($K$) of dynamic programming.

Q3: This is a fantastic question. The point to note is that the convergence bound for EWP is dimension-free, while those of the categorical algorithms are not (they depend on $d$). However, in the case of EWP-TD, the algorithm is prone to convergence to local minima (see our discussions on lines 146 and 324). Thus, we suspect that the EWP-TD examples are likely stuck in local minima, but notably that the quality of these local minima is roughly uniformly high for enough particles (e.g., as many as we use in our experiments). As you suggest, we would expect that EWP-TD should perform favorably to the Cat-TD algorithms with high $d$ -- however this is not guaranteed. Particularly, there is no known convergence guarantee for the EWP-TD algorithm when $d > 1$. Moreover, we could potentially leverage prior knowledge of the supports of the return distributions to enhance the resolution of the Cat-TD representations in relevant areas of the space of multi-returns, which could result in improved performance.

Q4: This was in reference to the Signed-Cat-TD algorithm. In previous experiments, the Cat-TD algorithms were representing categorical distributions on evenly-spaced points in the space of multi-returns, which necessarily requires an exponential number of support points as a function of $d$. On line 337, we meant to say that we instead fixed a number $m$ of support points and generated supports on $m$ randomly-chosen points in the space of multi-returns to avoid the exponential blowup (at the cost of lower resolution). Notably, previous theory and algorithms for categorical distributional RL do not accommodate such supports; this is a novel feature of our algorithms and theory.

Thank you again for your thorough review and your enthusiasm for our work. We appreciated your thoughtful questions and have addressed them in our rebuttal. We wanted to check if our responses have adequately clarified your questions. If so, we would be grateful if you could consider increasing your score. If there are any remaining points you would like us to clarify, we would be eager to discuss further.

Thanks to the reviewer for the questions. We really appreciate your engagement and your effort reading further into the proofs.

Proposition 1, Line 600. Thank you for pointing this out, we will clarify this. Since $\overline{\mathrm{MMD}}$ is a metric on multi-return distribution functions, we employ the usual notion of a ball in a metric space.

Line 617 and 621. Again, thank you, these are typos. Indeed on line 617 we will fix $\eta$ to $\eta_k$ and on line 621 we will fix $\eta$ to $\eta_K$.

Blowup of number of particles. There are two points to consider here. We can first ask how many large the representations become after a certain number of iterations. Additionally, we can ask, for a given error tolerance $\epsilon$, how large does $m$ need to be.

With regard to the first question, the number of particles in Theorem 3 remains fixed over the course of all iterations. With the unprojected DP algorithm, this number will increase at an exponential rate as discussed with 9dnf. You are correct that after a finite number of iterations, we will still have finitely many particles; see our discussion for the next point that shows that this number will still be much larger than what we get with Theorem 3.

The problem that Theorem 3 solves is the following: "Given a budget of $m$ particles, give me an $m$-particle EWP representation of that achieves error at most $\epsilon$". Theorem 3 says that, to accomplish this, you need

$

m_{\mathrm{ours}} \geq \widetilde{O}\left(\frac{1}{\epsilon^2}\frac{d^\alpha R^{2\alpha}_{\max}}{(1-\gamma^{\alpha/2})^2(1-\gamma)^{2\alpha}}\log^2\left(\frac{|\mathcal{X}|\delta^{-1}}{\log\gamma^{-\alpha/2}}\right)\right).

$

The number of iterates $K$ is an algorithmic detail for accomplishing this in our case -- we solve this problem in polynomial time ($K$ is polynomially large).

Suppose instead you try the method shown in the example of exponential blowup. You start with a 1-particle EWP representation. Note that the distributional Bellman operator is a $\gamma^{\alpha/2}$-contraction, so you'll have $\overline{\mathrm{MMD}}(\eta_k, \eta^\pi)\leq \gamma^{\alpha k/2}D$, where $D =\overline{\mathrm{MMD}}(\eta_0, \eta^\pi)$. So, if you want at most $\epsilon$ error, you need $K\geq \frac{2\log(D/\epsilon)}{\alpha\log\gamma^{-1}}$. Then, since each iteration blows up the number of particles by a factor of $|\mathcal{X}|$, this results in

$

m_{\mathrm{unprojected}} \geq |\mathcal{X}|^{\frac{2\log (D/\epsilon)}{\alpha\log\gamma^{-1}}}

$

To summarize, as a function of $\epsilon$ and $\mathcal{X}$, we have $m_{\mathrm{ours}} = \widetilde{O}(\log^2(\lvert\mathcal{X}\rvert)\epsilon^{-2})$ while $m_{\mathrm{unprojected}} = O(\lvert\mathcal{X}\rvert^2\epsilon^{-2})$, which is much worse.

The reason why we need to specify $K$ is because we're using randomized projections; applying too many raises the (low) probability of sampling a bad projection. We found the "just right" $K$ that avoids this with arbitrarily high probability. But again, $K$ is not to be interpreted as a user-chosen parameter here. It is an algorithmic quantity that gets us $\epsilon$-approximate return distributions for much smaller $m$ than required for unprojected DP.

Thank you for the author's response. However, after carefully reading the proof of Theorem 3, I'm not convinced by the meaning of the theorem and its proof.

1. The author claims that the free parameter is $m$, but this makes Theorem 3 seem to only propose an upper bound on MMD at a specific iteration $K = \lceil \frac{\log m}{2 \log \gamma^{-c/2}} \rceil$. In short, for a given $m$, the phrase "after enough $K$ iterations" does not seem valid. In my view, the theorem should be expressed to indicate that for any arbitrary $K$ and if there are $m = O(\exp{K})$ atoms, then an upper bound on MMD can be provided, and that as $K \rightarrow \infty$, the bound can be sufficiently reduced to approach zero.

2. The expression in Proposition 1, Line 600, seems mathematically awkward. Based on the results in Lines 599-600, the MMD as $k \rightarrow \infty$ implies $\frac{f(d,m)}{1- \gamma^{c/2}}$, but Line 600 expresses this in terms of a distribution and a ball, which has not been well-defined for MMD in this context. While this appears to be a minor expression issue that does not affect the result of Theorem 3, it still seems necessary to correct it.

3. It seems that $\eta$ on Line 617 should be changed to $\eta_k$, and $\eta$ on Line 621 should be changed to $\eta_K$.

After reading reviewer 9dnf's question and the author's answer, the author claims to have avoided an “exponential blowup”. However, at this point, I do not believe that Theorem 3 successfully solves this problem. I think more explanation is needed from the author, and thus I decreased the score to 5.

We thank the reviewer for their careful assessment of our work.

Could the authors comment on why they did not observe in practice the convergence suggested by their theorem in the EWP-based technique?

We did in fact observe convergence of the EWP-based algorithms in practice -- can you please point us more specifically to where the text or results suggest otherwise?

In Monte Carlo approximation, the price of increasing the dimensionality is often hidden in the “variance”...

This is a very interesting question! We suspect that it may be possible to provide "instance-dependent" bounds that sharpen with more favorable correlation structure between reward dimensions. Our results leverage the boundedness of the reward function and the structure of strongly negative-definite kernels to provide worst-case bounds. The polynomial scaling of the bound with dimensionality is inherited from the fact that the worst-case values taken on by the MMDs we consider also scale polynomially with dimension.

Regarding experimental results, please see our general response, in which we demonstrate that the signed categorical TD algorithm proposed in this paper is straightforward to combine with neural network function approximation.

Thank you again for your thorough review and your positive comments on our theoretical results. In our rebuttal, we addressed your questions and provided a demonstration of a larger scale application of our algorithm as per your suggestion. We hope these additional results have helped address your concerns about the scalability and practicality our work. We wanted to respectfully check if you've had a chance to review our rebuttal. We're keen to understand if the new illustration results and our responses to your questions have adequately addressed your concerns. If so, we would be grateful if you could consider increasing your score. If there are any remaining points you'd like us to clarify, we would be eager to discuss further.

We thank the reviewer for their feedback on the paper. We address the queries on notation below, and will include clarifications in the revised draft.

• The RHS of equation (2) is simply the set of all empirical distributions on $m$ points. That is, the distributions obtained by picking $m$ points $\theta_i\in\mathbf{R}^d$ and constructing the distributions supported on those points with equal mass. This is often referred to as a $m$-quantile representation in the $d=1$ case (see e.g. [BDR23]) and is equivalent to the model of [ZCZ+21].
• The object $\mathcal{R}$ in line 187 is defining a state-dependent support map. That is, for any input state $x$, $\mathcal{R}(x)$ outputs a finite set of support points $\{\xi_i(x)\}_{i=1}^{N(x)}\subset\mathbf{R}^d$. In the $d=1$ case, these would be the locations of the atoms (bins) of the categorical return distributions (see e.g. [BDM17b]) -- we generalize the notion here to have categorical supports that can depend on the state. The quantities $N(x)$ describe how many support points are in the categorical support of state $x$ -- indeed, we additionally generalize the model beyond that of [BDM17b] to accommodate state-dependent resolution of categorical supports. Thus, at state $x$, our model consists of $N(x)$ support points in $\mathbf{R}^d$, and these are indexed by $i$ in the equation in question.
• The $\Delta$ on line 188 represents a simplex. In particular, $\Delta_A$ for a finite set $A$ is the set of probability mass functions on $A$.

Thank you for pointing out these clarity issues, and we will clarify these in the revised draft. Should you have any other questions, we would be happy to discuss further.

Thank you for taking the time to review our submission. We appreciate your feedback regarding the notation used in our paper. We have addressed your questions in our rebuttal and hope that our clarifications have been helpful. If there are any other aspects that you would like further clarification on, we would be more than happy to discuss further; please let us know. If we have addressed all of your concerns, we would be grateful if you could consider increasing your score.

Once again, thanks for the efforts on rebuttal. I have adjusted my score. However, I do want to emphasize to the authors on the clarity of definitions, notations and motivations(in particular, categorical representation section) in the future revisions.