Paper: Distributional Bellman Operators over Mean Embeddings (ICLR-2024)

Reviewer comment (Novelty):
**Limited Methodology contribution and novelty**. Firstly, the sketch is just one special statistical functionals equipped with a specified kernel function, which is typically less commonly used than quantiles or samples. Working on algorithms simply based on a new concept may not contribute to the real development of this research field. More importantly, feature mapping has already been implicitly considered in MMDDRL paper [1], where MMD naturally induces a feature map. Notably. [1] also shows that not all kernels will lead to a convergent distributional Bellman operator, about which this paper has not fully discussed. From the perspective of new algorithms, Sketch-DP or TD is not new to me and even straightforward. Its concrete version, e.g., in Proposition 4.4 is very similar to the categorical representation. The choice of different feature maps also has a strong correlation with the MMD equipped with different kernels. Based on my knowledge of RKHS, different kernel naturally induces a RHKS with a specific feature map. Therefore, I do not think the sketch framework has sufficient methodological contribution compared with existing works.

Author response:
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.

Reviewer comment (Theory):
**Insufficient theoretical analysis**. The biggest issue in the theoretical part is the pre-specification of a metric \(d\) such that it is \(\gamma^c\) contractive, which is a strong assumption. With a metric that can already guarantee the contraction, the approximation error is easy to show. Further, as mentioned in the first weakness, not all feature maps can guarantee the contraction. For example, [1] gave some counterexamples for the convergence when we use Gaussian kernels. Therefore, the current paper lacks a crucial theoretical part about what kinds of feature maps can guarantee the contraction and then bound the errors. Without this crucial part, I personally think the current theoretical results are insufficient.

Author response:
**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.

Reviewer comment (Novelty):
**Insignificant empirical improvements**. The viewpoint of limited methodological contribution is further demonstrated by the insignificant empirical improvements in Figure 5. Although I truly agree on the conclusion that the sketch framework slightly improves C51 and QRDQN, it performs worse than IQN. Hence, I only view the sketch paper as a feature map-based algorithm, which has already been implicitly investigated before, rather than a novel one with significant performance. The insignificant improvement also let me rethink what is the real motivation of the sketch framework.

Author response:
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.

Reviewer comment (Experiments):
The authors' motivation for using the sketch Bellman operator is to reduce the need for expensive imputation strategies when converting between sketches and distributions. However, the experiment does not verify this claim. It would be helpful if the authors could provide some quantitative results (such as training time) to demonstrate this reduction.

Author response:
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.

Reviewer comment (Experiments):
The proposed method performs worse than IQN, even though the authors claim that IQN uses a more complex prediction network for non-parametric predictions. It would be beneficial if the authors could provide quantitative results (such as neural network sizes and training time) to explain this difference.

Author response:
**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.

Reviewer comment (Experiments):
It would be useful to include vanilla Statistical Functional Dynamic Programming (Bellemare et al., 2023) as a baseline in some toy examples to compare its performance with the proposed Sketch-DQN.

Author response:
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 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.

Reviewer comment (Writing):
The writings in some sections are not very clear. For example, in the subsection on 'Computing Bellman coefficients', the authors directly show the closed form of (5) without any explanations or citations.

Author response:
**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.

Reviewer comment (Writing):
There are some typos that need to be addressed, such as in Sec 2.2, where it should be $\left(\left(\mathcal{T}^\pi \iota (U)\right)(x)\right)$ instead of $\left(\left(\mathcal{T}^\pi \iota U\right)(x)\right)$.

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

Reviewer comment (Theory):
The proposed method requires a linear approximation for the Bellman update equation and requires calculating a Bellman coefficient matrix $B_r$ that can be computationally challenging. This contradicts the motivation to improve computation efficiency and reduce the imputation error by purely operating in the sketch space. It is unclear how the proposed method is superior to previous methods both computationally and statistically.

Author response:
**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.

Reviewer comment (Evaluation):
The experimental validation is limited. The experiment on the tabular tasks only shows the approximation error of the proposed sketch method rather than the performance of the overall approach. From the result, the proposed sketch method has a similar Cramer distance compared to the CDRL baseline. While the proposed sketch has lower excess Cramer distance and mean-embedding squared error, it is unclear how these metrics translate into the performance of the proposed approach.

Author response:
**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.

Reviewer comment (Evaluation):
For the Deep RL part, the proposed method underperforms IQN and does not significantly outperform QR-DQN, which is the backbone of the proposed method. Does the proposed method have a significant computation advantage?

Author response:
**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.

Reviewer comment (Presentation):
Minor: CDRL is mentioned but not referred to in the paper.

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

Reviewer comment (Theory):
What is the advantage of the proposed method compared with previous distributional RL methods? Does the proposed method have a provable lower approximation error or computation advantage?

Author response:
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.

Reviewer comment (Experiments):
How does the proposed method perform on the tabular tasks compared with baselines? Does the proposed method have a computation advantage in the deep RL setting?

Author response:
**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.

Reviewer comment (Theory):
The formulation of feature maps used in the paper (equation (8)) is sort of "pulled out a hat". It would be useful to justify why it made sense to use that formulation as opposed to other possibilities. See also the questions below.

Author response:
**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.

Reviewer comment (Theory):
The general sketch using Bellman coefficient $B_r$ will correspond to some mean embedding sketch which is an invertible linear combination of first $m$-moments, right?

Author response:
**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.

Reviewer comment (Theory):
Where's the generalization to the case when $\mathcal R$ is not finite?

Author response:
**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.