Distributional Bellman Operators over Mean Embeddings

We thank the reviewer for their time in reviewing the paper.

We respectfully disagree with the reviewer's assessment, and in particular we would like to clarify fundamental differences between our work and that of Nguyen-Tang et al. (2021), and clarify the validity and novelty of our theoretical results. We hope that our response below makes these points clear, and that the reviewer will consider revising their score; we would be very happy to engage in further discussion on any of these points if there is remaining uncertainty.

Methodology contribution, and novelty
The reviewer seems to claim that the methodology introduced in the paper is not novel, and is already encompassed by the work of Nguyen-Tang et al. (2021), stating "Sketch-DP or TD is not new to me". We strongly disagree with this, and we believe the works are fundamentally different, for the following reasons:

• The work of Nguyen-Tang et al. contributes an empirically highly successful approach to distributional temporal-difference learning. The core features of their approach are:
  • The approach approximates return distributions by learning particle locations, as with QR-DQN.
  • These particle locations are learnt through temporal-difference learning, by combining an MMD loss with target distributions formed by distributional bootstrapping.
  • As mentioned above, this approach is highly successful empirically, but the TD algorithm itself is not analyzed theoretically. The authors prove contractivity of the distributional Bellman operator with respect to certain MMDs, but this is typically only one aspect of many required to analyze a distributional RL algorithm, since the distribution approximation and TD dynamics must also be taken into account.
• In contrast, this work contributes a family of dynamic programming and temporal-difference learning algorithms for learning sketch values directly. Contrasting the characteristics of our work with those of Nguyen-Tang et al.'s work, we note:
  • Our approach learns sketch values directly, and does not approximate probability distributions through particle representations.
  • We derive a dynamic programming algorithm for updating sketch values based on the notion of Bellman coefficients, defined by the regression problem in Eqn (5). This results in a DP procedure expressible through straightforward matrix-vector products. We also derive an associated TD procedure through stochastic approximation of the DP procedure.
  • We provide a novel theoretical analysis of the complete DP procedure described above, based on the expression of the DP procedure via Bellman coefficients.

Since the work of Nguyen-Tang et al. (2021) and our work therefore take quite different approaches, we do not believe that the claim of lack of novelty of our work is supported. If the reviewer believes our method has appeared elsewhere, we would be happy to take into consideration further references; otherwise, we would be grateful to the reviewer if they could update their score accordingly.

Theoretical analysis
We believe the reviewer has misunderstood the requirements on the metric $d$. Proposition 3.2 requires that the usual distributional Bellman operator $\mathcal{T}^\pi$ is a contraction with respect to $d$. This is not a strong assumption at all, as standard metrics such as supermum-Wasserstein-$p$ metrics and the supermum-Cramer distance satisfy this requirement.

In addition, we believe there is a high degree of novelty in our theoretical analysis. As mentioned above, we give a complete analysis of a distributional RL algorithm, Sketch-DP, in contrast to analyzing the contractivity of the distributional Bellman operator, which is not sufficient in itself to analyze the behavior of practical distributional RL algorithms, due to not incorporating distribution approximation errors into the analysis (as in Nguyen-Tang et al). This analysis goes beyond the typical contractive mapping approach, and we believe this form of analysis will be useful in future theoretical research on distributional RL.

We agree that the performance of Sketch-DQN on the Atari suite is not superior to all existing distributional methods. We would argue that state-of-the-art performance in the Atari suite should not be necessary for acceptance, given that the paper primarily is concerned with contributing novel algorithmic ideas and theoretical analyses in distributional RL. These deep RL experiments show that our theoretically motivated approach is able to achieve good performance in Atari, achieving similar benefits as observed with earlier approaches to distributional RL.

"why bother to use an algorithm among two similar ones that is infererior in performance?"
The reviewer questions the value of having a new algorithm that performs better than all but one benchmark. First, we do not agree with the premise that our algorithm is similar to that of Nguyen-Tang et al. beyond superficial relations in RKHS (see above). Second, we do not intend to convince people to commit to our method a priori, but to show that this method works in combination with deep RL in practice and reaches or surpasses established benchmarks, under the same or similar computational budget, tasks and training procedures.

Distributional RL methods are used in a wide variety of applications, and at a wide variety of scales; whether this method will be superior or inferior to another in any given application depends on a wide variety of factors. The point is that we offered a novel family of DRL algorithms, notably with convergence analysis in the DP setting, with promising results, and hence it is worth considering for other applications.

Thank you for the additional comments.

The reviewer states that our description of MMDRL above is inaccurate, although we can't see any substantive inaccuracies highlighted in their response. In response to the two bullet points the reviewer raised:

• Both MMDRL and QR-DQN learn approximate distributions of the form $\sum_{i=1}^m \frac{1}{m} \delta_{z_i}$, and in particular learn the particle locations $(z_i)_{i=1}^m$. For QR-DQN specifically, the particle locations are trained to represent (approximate) quantile locations.
  • On methodology, the review asserts that a sample representation is advantageous over pre-defined statistics. To our knowledge, this isn't a uniform principle; whether each representation is better than another depends on many factors, and on the application. The reference [2] in the reviewer's response is not included; we would be happy to comment further if the reviewer provides this reference.
• The reviewer is correct that Theorem 3 and Proposition 2 in Nguyen-Tang et al. analyze approximation error of a fixed distribution with a set of particles under MMD. However, these results, and the contraction analysis of $\mathcal{T}^\pi$ in Theorem 2 of that paper, do not prove convergence of the algorithm proposed in that paper; based on this, we don't agree with the reviewer's claim of "lower[ing] their theoretical contribution". Please see the response on theoretical analysis for further comments on this.

“the largest concern is the overlap between the paper and [1] as both are RKHS-based methods”
We agree with the reviewer that both methods are related to the mathematical concept of RKHS, but we don't see substantial overlap beyond this, with notable variation along the dimensions of particle/sketch approximation, TD/DP algorithms, and theoretical convergence analysis in the case of Sketch-DP, as described in our first response.

We would like to emphasize that we view the work of Nguyen-Tang et al. as an important contribution in the field of distributional RL, and we believe this paper and Nguyen-Tang et al. make complementary, distinct contributions. To make this clearer, we have added Appendix B.7, which summarizes these points of comparison to existing approaches, as well as our response on the point of theoretical analysis below. We have also added a comment regarding the use of MMD by Nguyen-Tang et al. in the related work section.

"but the mathematical foundation especially the contraction would be very linked".
We agree with the reviewer that both approaches are related to the mathematical concept of RKHS. However, as mentioned above, we don't think there is a mathematical link between the two works beyond this, and as far as we can tell, the reviewer is not suggesting a concrete link in their response. If the reviewer has a concrete link between mathematical foundations in mind, we would be happy to discuss further.

We also believe the mention of "contraction" here is related to a misunderstanding regarding the relevance of showing contractivity of $\mathcal{T}^\pi$ with respect to certain metrics, described in more detail in the theoretical analysis section of the response below.

"This is also the major deficiency of CDRL compared with QRDQN. However, it seems that your proposed algorithm also suffers from this fundamental issue."
The reviewer asserts that the specific instance of the sketch algorithm in Prop 4.4 is similar to CDRL. However, CDRL uses triangular features, whereas this instance uses square indicator functions. The analysis of CDRL relies on the projection of measures on the Cramer distance, facilitated by the triangular features, while our analysis is motivated by bounded error analysis based on the Wasserstein distance. Therefore, CDRL and this method differ both practically and theoretically.

We also disagree that selecting distribution support in advance is a fundamental issue. Return distributions are very often bounded in applications of interest, which leads to a natural choice of return range. This idea is used in several modern high-performing RL agents, including MuZero (Schrittwieser et al., 2020), and DreamerV3 (Hafner et al., 2023).

References

Schrittwieser, Julian, et al. "Mastering atari, go, chess and shogi by planning with a learned model." Nature 588.7839 (2020): 604-609.

Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).

We believe there may be a misunderstanding here about what is required in the analysis of a distributional DP algorithm. We hope the explanation below makes clear that the paper contributes a full rigorous analysis for the introduced class of DP algorithms, and we also highlight that complete analysis of dynamic programming procedures in distributional RL, as in this paper, is a substantial theoretical contribution (see e.g. Rowland et al., 2023, Wu et al., 2023, for other examples). We would be very happy to continue the discussion if the reviewer has further questions.

"Linked with 1., a natural question is how you analyze the contraction, regardless of DP or TD, in terms of the kernel."
We believe the reviewer is asking about the contraction specifically of the distributional projection operator $\mathcal{T}^\pi$ (without any projection), with respect to the (supremum-)MMD induced by the choice of feature map.

The central point we would like to make is:
Establishing contractivity of $\mathcal{T}^\pi$ with respect to this (supremum-)MMD does not prove convergence of the algorithms we propose in the paper. Conversely, our convergence analysis does not make use of such a result.

Thus, contractivity of $\mathcal{T}^\pi$ with respect to the (supremum-)MMD is not in and of itself a relevant theoretical result for analyzing the algorithms proposed in the paper. Therefore, our response to the following point from the reviewer:

"However, I did not find the relevant theoretical results in the paper, while the only thing I found is the $\gamma^c$ contraction assumption."

is that it is not a relevant theoretical result, and our convergence analysis studies the DP algorithm without using such a result, with Proposition 4.4 using supremum-Wasserstein distance, for example.

"I am aware that this assumption is true for Wasserstein and Cramer distance, but it is not true for general kernel-based algorithms, which are highly linked with the choice of kernel. "
Given the explanation above, we expect this comment is related to the difference between:

• A metric used in defining a distributional RL algorithm.
• A metric used to analyze this algorithm.

We believe our discussion above addresses this concern. Our analysis does not require us to use the same metric used in defining the algorithm to analyze the algorithm.

"... I acknowledge the contribution of the following analysis. It is no doubt non-trivial to have such a detailed answer instead of directly imposing an assumption in the paper."
Thank you for this comment. We hope that with the explanations given above, there are no remaining theoretical concerns that the reviewer has, but please let us know if not.

Further details on the role of contraction in distributional RL convergence results
To provide additional context as to why it is not necessary to prove convergence of $\mathcal{T}^\pi$ with respect to a certain metric to establish convergence results for distributional dynamic programming algorithms, we recall some examples from the literature:

• In the case of quantile dynamic programming (see Rowland et al., 2023, and also Dabney et al., 2018 for earlier analysis), the algorithm is defined with a Wasserstein-1 projection, but the convergence analysis makes use of a result showing contractivity of $\mathcal{T}^\pi$ with respect to the supremum-Wasserstein-$\infty$ metric. This contractivity result is then combined with a non-expansion result for the projection mapping component of quantile-dynamic programming to prove convergence.
• In the case of fitted likelihood estimation (FLE; Wu et al., 2023), one result that contributes to the analysis of the main algorithm in the paper is the contractivity of $\mathcal{T}^\pi$ with respect to a weighted sum of Wasserstein-2 distances at each state, but note that the Wasserstein-2 distance is unrelated to the proposed algorithm itself. This result is then combined with an error propagation analysis to obtain a proof of convergence.

In contrast, the work of Nguyen-Tang et al.: (i) establishes contractivity of $\mathcal{T}^\pi$ with respect to certain supremum-MMD metrics, (ii) analyzes approximation of fixed distributions in MMD. However, these results are not used to prove convergence of the algorithm proposed in the paper, and the algorithm proposed in the paper is not analyzed. Analysis of their proposed TD algorithm would require analysis of the dynamical system induced by the incremental particle updates; see Rowland et al. (2023) for an example in the case of quantile TD learning.

We thank the reviewer for their time and constructive feedback in reviewing our paper. We respond to specific queries raised by the reviewer below. We are eager to engage in further discussion on any of these points if the reviewer has further queries, and hope that the reviewer will consider re-evaluating their score as a result.

Comparison with statistical functional dynamic programming
Based on your suggestion, we ran a version of statistical functional dynamic programming (SFDP) using expectiles and the imputation strategy described by Bellemare et al. (2023) for comparison against Sketch-DP. We use SciPy's default minimization algorithm to compute the imputation strategies required by SFDP. We have included the full results in Appendix D.4 in the latest version of the paper. A brief summary is provided below to provide quick answers to the reviewer's points 1 & 3.

• Point 1. Running SFDP with expectiles, and Sketch-DP with sigmoid features in the environments from the paper, we found that the per-update wallclock time for SFDP was consistently greater than 100x the wallclock time for Sketch-DP. This is due to Sketch-DP requiring only matrix-vector product computations, while SFDP requires the solution of an optimization problem in each call it makes to the imputation strategy.
• Point 3. With the same experimental set-up, we found that the reconstruction errors (measured in Cramer distance) for Sketch-DP and SFDP to be comparable for low numbers of features/expectiles, with some advantage to the Sketch-DP approach being observed for higher numbers of features/expectiles.

Thus, these results confirm that there are significant practical speed-ups that come from the Sketch-DP approach relative to SFDP. We thank the reviewer for suggesting this additional comparison; full results are given in Appendix D.4, please let us know if you have any further questions on these areas.

IQN training times. Based on the suggestions of several reviewers, we reported measured wallclock times for Sketch-DQN, QR-DQN, and IQN in the general response above, to support the discussion around the additional complexity associated with IQN in the paper.

Derivation of Equation (7). The form for the Bellman coefficients in Eqn (7) follows since Eqn (5) defines a least-squares linear regression problem, which B_r solves. The expression $C_r C^{-1}$ is the analogue of the usual $\hat{\beta} = (X^T X)^{-1} X^T Y$ expression for the least-squares coefficients for the linear model expressed in the form $Y = X \beta + \epsilon$. We thank the reviewer for raising this question, and agree that this derivation can be made clearer; we have added a guide to the derivation in Appendix B.6, and in doing so, corrected a typo in Eqn (7). Please let us know if you have any further queries.

Typo. Thank you for noting this, and please let us know if there are any other typos you had in mind for correction.

We thank the reviewer for their time and constructive feedback in reviewing our paper. We respond to specific queries raised by the reviewer below. We are eager to engage in further discussion on any of these points if the reviewer has further queries, and hope that the reviewer will consider re-evaluating their score as a result.

Computation of Bellman coefficients
The computation time of the Bellman coefficients $B_r$ is insignificant, particularly in the setting of deep reinforcement learning, since it is a one-off computation done before running the recursive updates; see Algorithm 1, where $B_r$ is computed before the main loop. For example, it takes only a couple of seconds on a CPU to compute $B_r$ for the case with 100 feature functions, and $n=10^5$ grid points to estimate the expectations in Eqn (7). This is negligible in the Atari experiments, where only rewards of -1, 0 and 1 are encountered, relative to the many forward/backward passes made in training deep networks. We have also added additional experiments comparing the computation cost associated with Sketch-DP with an earlier distributional RL algorithms (SFDP) in the tabular setting in Appendix D.4 (please see the response to Reviewer RJqf for further details), which confirm that Sketch-DP delivers computational improvements relative to SFDP in the tabular setting too.

Additionally, in settings where possible reward values are not available in advance, dealing with unseen reward values $r$ does not require solving the regression problem from scratch. This is because $C^{-1}$ can be saved, and we only need to compute $C_r$ for unseen values of $r$, which is also very fast; we have added a paragraph to Appendix B.6 to describe the details of this approach.

Comparison against previous methods in combination with deep reinforcement learning
Based on your suggestion, we have reported new results in the general response above, comparing wallclock rates of frame throughput for Sketch-DQN, QR-DQN, and IQN. As described in detail in the general response, Sketch-DQN attains a higher throughput than QR-DQN and IQN due to the straightforward nature of the distributional update derived for sketches in this paper. We thank the reviewer again for this suggestion. We would also like to emphasize that although QR-DQN and Sketch-DQN share the same network architecture, the two agents use completely different losses for updating their predictions.

Advantage over previous methods. We would like to emphasize that some of the core advantages of the sketch framework introduced in this paper are (i) the generality of the framework - each choice of feature functions yields a new algorithm; and (ii) the way in which these methods can be analyzed, as described in Section 4. We expect these contributions to be of use in future distributional RL work, beyond the current paper. Our experimental results also show benefits of Sketch-DP/TD algorithms, but we emphasize that the contributions of the paper extend beyond the empirical performance of the newly introduced algorithms.

Focus on approximation error in tabular experiments. We believe the reviewer is highlighting the difference in control and prediction tasks within reinforcement learning; our deep RL experiments are focused on measuring the performance of Sketch-DQN in control, and in our tabular experiments, we focus on understanding the quality of predictions made by the sketch algorithms introduced in the paper. Measuring approximation error is an important aspect of empirical evaluation of the methods; it allows us to understand how well a particular algorithm approximates the return distribution in various settings, and from this perspective we view it as an important complement to the deep RL experimental results. Further, it has been shown that in tabular settings, distributional reinforcement learning may not necessarily offer improved performance for control (Lyle et al., 2019), meaning that in the tabular setting, evaluation of distribution predictions is often more informative as a means of understanding the algorithm's properties.

Minor points. Thank you for mentioning the CDRL acronym, we will ensure it is defined and referenced in the final version.

References

Clare Lyle, Pablo Samuel Castro, Marc G. Bellemare, A Comparative Analysis of Expected and Distributional Reinforcement Learning, AAAI Conference on Artificial Intelligence, 2019.

We thank the reviewer for their time and constructive feedback on the paper, and are glad to hear they enjoyed reading the paper so much! Please see below for responses on specific queries.

Formulation of translation-family feature maps. The translation family of features is a natural choice because it is simple to implement and produces good empirical results (as reported in the experimental results in the paper). These functions are similar to radial basis functions used elsewhere; in particular, in prior work on distributed distributional codes (Sahani & Dayan 2003, Vértes and Sahan 2018, 2019). We also emphasize that other families of feature maps can also be considered, and an empirical comparison across a wider variety of feature maps would be an interesting direction for future work.

The general sketch and first $m$-moments
We are unsure we understand the question fully, so have given a summary of the theory relating to mean embedding sketches and when they can be related to first $m$-moments below; please let us know if we have misunderstood the query or if you have further questions, however.

For a given collection of feature functions $\phi$, and corresponding Bellman coefficients $B_r$, the corresponding computed sketch values do not necessarily correspond to invertible linear combinations of moments. However, the existing theory referenced (Theorem 4.3 of Rowland et al., 2019), guarantees that if the regression error incurred in Equation (5), defining the Bellman coefficients, is 0, then the sketch must correspond to an invertible linear combination of the first $m$-moments, as the reviewer states. We hope this answers the question, but please let us know if not.

Why can we write $B_r$ as $C_rC^{-1}$? This follows since Equation (5) defines a least-squares linear regression problem, which $B_r$ solves. The expression $C_r C^{-1}$ is the analogue of the usual $\hat{\beta} = (X^T X)^{-1} X^T Y$ expression for the least-squares coefficients for the linear model expressed in the form $Y = X \beta + \epsilon$. We thank the reviewer for raising this question, and agree that this derivation can be made clearer; we have added a guide to the derivation to Appendix B.6, and in doing so, corrected a typo in Eqn (7).

Arbitrary reward distributions
Thank you for raising this: The generalization to arbitrary reward distributions is obtained by modifying the term $\phi(r+\gamma G)$ in the definition of $C_r$ in Eqn (7) to $\mathbb{E}_{R \sim \rho} [\phi(R + \gamma G)]$, where $\rho$ is the distribution of the immediate reward. The intuition behind this is that the regression target is now the average of the sketch values, over the possible values of the immediate reward. To obtain an algorithmic implementation, this is a one-dimensional integral which can be approximated with standard numerical integration techniques, as mentioned in Appendix C.1 (Sketch DP under conditional independence). Experiments on this version of the sketch Bellman operator are shown in the DC+Gaussian R columns of Figure 4, and more thorough evaluations are in panels c and d of Figure 7, and panel b of Figure 8.

Additionally, in settings where possible reward values are not available in advance, dealing with unseen reward values $r$ does not require solving the regression problem from scratch. This is because $C^{-1}$ can be saved, and we only need to compute $C_r$ for unseen values of $r$, which is also very fast; we have added a paragraph to Appendix B.6 to describe the details of this approach.