Paper: Diverse Projection Ensembles for Distributional Reinforcement Learning (ICLR-2024)

Reviewer comment (Novelty):
The description of the motivation and related work appears to be rather insufficient.

Author response:
We thank you for this suggestion and have expanded the related work section: specifically, we added references to Lindenberg et al.[1], Jiang et al.[2], Clement et al.[3], and Parisi et al.[4] with extended discussions. We also made adjustments to our motivating section in the introduction to improve clarity.

Reviewer comment (Presentation):
In Section 6.2, “the bsuite consists of 22 tasks” seems incorrect as Fig 9 shows 23 bsuite tasks.

Author response:
We corrected the Section 6.2 to say 23 tasks.

Reviewer comment (Reproducibility):
Appendix B.1, Table 4: DLTV has a fixed decaying schedule, $\sqrt{\logt /t}$. How can $\beta_{final}$ be defined?

Author response:
Indeed, $\beta_{final}$ should be $n/a$ in this entry. We furthermore removed the additional row $\beta_{init}$.

Reviewer comment (Presentation):
Appendix B.2, Table 5: Initial bonus $\beta_{init}$ is written twice.

Author response:
Indeed, $\beta_{final}$ should be $n/a$ in this entry. We furthermore removed the additional row $\beta_{init}$.

Reviewer comment (Writing):
Appendix B.2, Equation (17): What does $\eta_{M,\phi}$ refer to? The expression seems to be ambiguous.

Author response:
$\eta_{M,\phi}$ is a distributional model equivalent in architecture to $\eta_{M,\theta}$ but is trained on intrinsic rewards instead of extrinsic rewards with the goal of estimating the bonus $b_\phi(s,a)$ in the spirit of Theorem 3. Note that we do not require this estimate to be a distributional model per se, but found it convenient to implement given that it equals our procedure for extrinsic return estimation.

Reviewer comment (Theory):
In the paper, an upper bound is obtained by Proposition 2, but it seems insufficient to justify the proposed bonus $b_{\phi}$ as an epistemic uncertainty estimate. Since the confidence bound in Proposition 2 is just an expression naturally induced by the 1-Wasserstein distance, it seems invalid from a regret analysis perspective. Furthermore, the authors again approximate the total distributional error by the one-step TD distributional error, which seems to have no significant theoretical connection. Of course, I think it is quite difficult to describe this theoretically, so I suggest to plot the experimental results of the bonus estimate $\hat{b}_{k+1}(s,a)$ or $\mathbb{E}_{s,a}[\hat{b}_{k+1}(s,a)]$ in each environment.

Author response:
We thank you for your question and suggestion to plot the bonuses $b_\phi(s,a)$. To answer your question: Our theoretical results do indeed not merit a full-fledged regret anlysis. We consider this a tall order even for many basic tabular algorithms. However, Proposition 2 proposes a bonus that links total distributional errors to optimistic overestimation of Q-Values, the basis for many furthergoing regret analyses. The connection of this bonus (the total distributional error) to the one-step TD errors is given by the decomposition according to Theorem 3. To illustrate the behavior of this bonus estimate in practice, we have included a Figure in the appendix that plots the bonus estimates $b_\phi(s,a)$ compared to inverse counts and the local error estimates for the entire state-space over several timesteps in the deep sea environment.

Reviewer comment (Experiments):
In Figure 4, is PE-DQN [Ind.] without uncertainty propagation equivalent to the C51/QR mixture model with greedy action selection? If so, the performance improvement of PE-DQN seems to be mainly due to the introduction of bonus estimates. Based on the experimental results in the current paper, it is difficult to confirm whether the elimination of the inductive bias by projection contributes to the performance improvement.

Author response:
The referenced section could indeed be clearer and we updated the writing in the current revision. PE-DQN[Ind.] is not equivalent to a C51/QR-DQN mixture model with greedy action selection, but with an optimistic action-selection. I.e., the agent adds a bonus $\beta w_1(\eta_{C51}, \eta_{QR})(s,a)$ to the action-selection, where $\eta_{C51}$ and $\eta_{QR}$ have been trained completely independently, obviating an additional bonus estimation procedure as implied by Theorem 3. Our further ablation experiments, PE-DQN[C51/C51] and PE-DQN[QR/QR], address your second concern of whether the improved exploration performance is in fact due to the diversity of projections in our ensemble, or solely due to the bonus estimation procedure. In this ablation, both PE-DQN[C51/C51] and PE-DQN[QR/QR] make use of the same bonus estimation procedure as PE-DQN, except that ensemble members are based on the same projection and representation.

Reviewer comment (Novelty):
Projection ensembles seem to be an attempt to relax the priors of the model as in IQN[1] and FQF[2]. Can the authors explain the difference from those?

Author response:
As you correctly point out, both IQN and FQF, as well as our model, have a higher representational capacity than QR-DQN or C51 in terms of the representable distributions. However, our main focus lies on exploiting the difference in inductive biases and the subsequent difference in generalization behavior, to more reliably quantify epistemic uncertainty in unseen $(s,a)$-regions. Neither IQN nor FQF have this functionality built in and accordingly neither can make use of this in an exploration setting.

Reviewer comment (Experiments):
since the proposed algorithm is an ensemble on distributional RL, comparisons only with either distributional RL or ensemble on scalar RL but not both seem inadquate. Perhaps worth comparing to some of the works you mentioned in Related Work.

Author response:
**1. Experimental baselines**: We agree that additional baselines, especially ones that leverage distributional ensembles for exploration would improve the comparative evaluation. The works we reference in Section 2 (e.g., Eriksson et al. [1] and Hoel et al. [2]) are not aimed specifically at exploration and would thus require additional adjustments to make for a sensible comparison in our context. However, other reviewers have pointed towards the recent work by Jiang et al.[3], which should indeed be a sensible comparison. We may consider adding this to the final version, but are afraid the current revision does not allow this due to time constraints.

Reviewer comment (Presentation):
the illustration for Fig.1 and Fig.2 can be expanded

Author response:
**2. Descriptions of Fig. 1 and Fig. 2**: This is a good suggestion and we extended the descriptions for both figures.

Reviewer comment (Theory):
the atoms in a quantile representation can be thought of with equal weights whereas those in a categorical representation cannot. Would you not need to make this distinction in Eq. 7?

Author response:
**2. Eq. 7**: The distinction between the equal weighting of Diracs in a quantile representation and the varying weights of Diracs in the categorical representation is implicitly contained in $\Pi_i \eta(s,a)$. For the quantile projection, $\Pi_q \eta(s,a)$ yields a mixture of equally weighted Diracs, whereas the categorical projection $\Pi_c \eta(s,a)$ yields a non-uniformly weighted mixture of fixed-location Diracs.

Reviewer comment (Theory):
the local estimate of the exploration bonus $w_1 (\hat{\eta}, \Omega_M\mathcal{T}^\pi\hat{\eta})$ seems to be measuring the distance between the ensemble return distribution as a whole and its backup target, I failed to see how it may be estimated as the model disagreement among ensemble members (page 7)

Author response:
**3. Local error estimation**: Informally speaking, this procedure is motivated by the fact that each ensemble member in $\eta_{M,\theta}(s,a)$ is trained to regress against the backup target $\Omega_M \hat{\mathcal{T}}^\pi \eta_{M,\theta}(s,a)$. The underlying notion being, that the disagreement of a deep ensemble provides a useful estimate of the expected error between individual predictions and the true labels. This is not a rigorous argument, but rather motivated through the many empirical results of deep ensembles (see Lakshminarayanan et al.[4], Riquelme et al.[5], or Fort et al.[6]), where some works aim to make such statements more formal through connections to the Bayesian setting (e.g., Osband et al.[7], B. He et al.[8], Izmailov et al.[9]). In the current revision, Appendix B.3 contains an illustration of the types of uncertainties used in PE-DQN in practice.

Reviewer comment (Experiments):
The performance improvement over existing methods is relatively small. It would be perhaps ideal to find scenarios where the performance of the PE-DQN is more significant than existing methods. One task that I can think of is Procgen. Recently, [1] showed that incorporating epistemic uncertainty can significantly improve the performance of Value-based methods on Procgen but the method in [1] does not incorporate future uncertainty and only uses an ensemble of QR-DQNs. I would expect PE-DQN to be much better in this case.

Author response:
We agree that it would be very interesting to evaluate PE-DQN on tasks like procgen and that a sensitivity analysis of $\beta$ would be a useful addition. For the current revision this unfortunately seems out of scope due to the computational burden in the limited time.

Reviewer comment (Evaluation):
I would like to see a comparison of the wall clock time to get a better understanding of the computational cost.

Author response:
We thank you for the useful suggestion and included a wall-clock comparison table for the computationally more demanding VizDoom experiments in the appendix.

Reviewer comment (Experiments):
I wonder if using different architectures or different ensemble weights for the ensemble would further improve the performance. Perhaps there is even a way where one can adjust the weights to get a tighter estimate of the uncertainty.

Author response:
**Adjusting weights**: Indeed, this is an interesting avenue for future work. For example, one may define an operator that is "greedy" in its model weighting w.r.t. some criterion (e.g., TD loss of each model) to obtain a mixture that achieves a "best of both worlds" behavior. Similarly, as you suggest, one may consider weighing models to obtain a tighter uncertainty estimate. Intuitively, this should make sense in situations where some models are significantly less reliable than others, where an equal weighting may be overly conservative. The central challenge in this, we suspect, is to find a well-calibrated criterion that does not sacrifice the reliability of uncertainty estimates.

Reviewer comment (Experiments):
How does this propagation scheme compare to other schemes like UBE [4]?

Author response:
**Comparison to UBE**: There are a few conceptual differences between our propagation scheme and the uncertainty Bellman equation (UBE) (O'Donoghue et al.), with the clearest being that our bound addresses uncertainty in a distributional sense, whereas the UBE quantifies uncertainty in expected returns. Moreover, our propagation scheme is motivated by a more model-free derivation that directly decomposes total (distributional) errors into distributional one-step TD errors, whereas O'Donoghue et al.[4] decompose the posterior variance of value estimates into posterior variances of reward and transition models. Nonetheless, there are several similarities and perhaps a Bayesian treatment of distributional value estimation may make the links and differences clearer, for example it could be interesting to analyze the tightness of the underlying bounds for these propagation schemes.

Reviewer comment (Writing):
There are many places where imprecise writing makes the presentation difficult to parse. For example, in the abstract it is stated that "profounding impacting the _generalization behavior_ of learned models". What does generalization behavior here mean? It is known that generalization in the context of RL can have many different meanings. As another example, in the introduction it is stated that "the space of potential return distributions is infinite-dimensional". The space of probability measures (as opposed to finite signed measures) is _not_ a vector space, and so how do the authors define a dimension in this context? It couldn't be the cardinality of the set because then the set of representable distributions will also have infinite dimensions. Or do the authors mean parametric vs nonparametric, where the dimension is defined on the parameter space? As another example, Figure 1 is poorly explained — what are the x and y axis, how does it convey the notion of "distinct generalization signature" as stated in the paper? As another example, in section 3, the notation is inconsistent. $S_t$ and $A_t$ are used to denote samples or random variables? If samples, then it is inconsistent with what is stated next that capital letters denote random variables. If random variables, then what does $\mathcal{R}(\cdot \mid S_t, A_t)$ or $\mathcal{R}(S_t, A_t)$ (two ways of stating the same thing?) even mean? $\mathcal{D}^\pi(s,a)$ is in $\mathscr{P}(\mathbb R)$ and not in $\mathscr{P}(\mathbb R)^{\mathcal S \times \mathcal A}$. In the definition of a pushforward measure in the appendix, it is stated that $B \in \mathbb R$, which should be $B \in \mathcal{B}(\mathbb R)$ where $\mathcal{B}(\mathbb R)$ denotes the Borel subsets of $\mathbb R$.

Author response:
**Generalization** We acknowledge that the term "generalization" is not precise and we adjusted our wording in the revision. Intuitively, we considered the term "generalization behavior" to mean a model's predictions on data-points not contained in the dataset. This coincides closely with the description by Kirk et al.[2] for online learning in singleton MDPs (see p. 203 ¶ Scope). If models of an ensemble generalize differently, ensemble disagreement can then be used to quantify epistemic uncertainty in novel inputs.

**"the space of [...] probability distributions is infinite-dimensional"**: This phrase is indeed does not contribute to clarity and we accordingly replaced it in the revised version. Our notion behind this statement was that one may aim to characterize a probability distribution, for example, through the sequence of its moments. Here, an infinite-dimensional vector containing moments may characterize almost arbitrary distributions. Still, even in this kind of embedding further restrictions would be required, hence our removal of the phrase.

**Description Fig. 1**: We agree that the description of Fig. 1 benefits from an extension and changed it.

**Corrections and Notation**: We thank you for pointing out several issues, most of which we have incorporated into the draft: The expression $\mathcal{R}(S_t, A_t)$ should indeed be $\mathcal{R}(\cdot|S_t, A_t)$. $S_t$ and $A_t$ indicate variables as stated and $\mathcal{R}(\cdot|S_t, A_t)$ denotes the reward distribution conditioned on the random variables $S_t$ and $A_t$ as opposed to the explicit events $S_t=s, A_t=a$ where $s$ and $a$ are known to have occured. We furthermore adjusted $\mathcal{D}(Z^\pi(s,a))\in \mathscr{P}(\mathbb{R}^{\mathcal{S}\times\mathcal{A}})$ to $\mathcal{D}(Z^\pi(s,a))\in \mathscr{P}(\mathbb{R})$ and updated the definition of the pushforward measure to read $B \in \mathcal{B}(\mathbb{R})$ where $\mathcal{B}(\mathbb{R})$ are the Borel subsets of $\mathbb{R}$, as suggested. You are correct that $p \in [1,\infty)$ for the $p$-Wasserstein distances and we changed this accordingly. We changed the passage "$\eta_M$ [...] has support over [...]" to express the set of representable mixtures $\mathscr{F}_E = \{ \eta_E(s,a) \,|\, \eta_E(s,a) = \frac{1}{M} \sum_i \eta_i(s,a), \, \eta_i(s,a) \in \mathscr{F}_i, \, i\in[1,...,M]\}$. We respectfully disagree with your counterexample, as $\mathscr{F}_1$ and $\mathscr{F}_2$ are defined to be sets of representable distributions, not random variables. The mixture distribution then is the convex combination of two normal distributions rather than the distribution of the sum of the random variables as implied in your example.

Reviewer comment (Theory):
Apart from the multiple questions asked in the weaknesses section, what is the use of Theorems and Propositions in Section 4, when they play no role in the sections ahead? If they are not empirically useful results, but simply theoretical curiosities, it is still unclear what is learned from them. They seem like simple manipulations of the various definitions without an end goal in sight.

Author response:
**Clarification of propositions and theorems in Section 4**: We would like to clarify the implication of our results in Section 4. Based on the notion that the choice of projection influences the generalization behavior (in the above-described sense) of a distributional model, we hypothesize, an ensemble that combines different projections can more reliably quantify epistemic uncertainty in unseen $(s,a)$ regions through the disagreement of ensemble members. However, even under the assumption that ensemble disagreement estimates errors between predictions and true targets, disagreement would not yield a useful exploration bonus since we use TD targets not the true return distribution (a known issue addressed for example by Flennerhag et al.[3], O'Donoghue et al.[4], and Fellows et al.[5]). Instead, a bonus that would allow for overestimation of $Q$-values is stated in Proposition 2 and can be decomposed into more accessible distributional TD errors according to Theorem 3. We respectfully disagree that these results are not empirically useful, as they are explicity used in our algorithm and showed consequential in ablation studies in Fig. 4b) rhs. We have also included illustrations in the appendix B.3 that compare these types of uncertainty.