A Principled Path to Fitted Distributional Evaluation

Thank you very much for your efforts. Please find our responses to your comments below.

Weakness-1, Question-1,2: Relaxing the strong assumptions

We acknowledge that both completeness (Assumption 3.1) and data splitting are limitations of the current work. While these are substantial directions that cannot be handled by the current work, we agree these limitations are worthy to be addressed. However, it is also true that these are standard limitations of many iteration-based RL algorithms. Many papers assume completeness (e.g., papers listed in Lines 247-248) and adopt similar strategies to data splitting or equivalently fresh samples in each iteration (e.g., [1], [2]) for theoretical convenience.

Answer to your Question 1: When Assumption 3.1 (Bellman completeness) fails, we will have a non-negligible approximation error, similar to FQE, at each iteration of our proposed FDE algorithm. Such an approximation error can be quantified by inherent distributional Bellman error, which can be similarly defined by adapting from the standard notion in FQE. See [3] and [4] for more details.

Answer to your Question 2: Thanks for your suggestions. Due to the general choice of functional Bregman divergences, conducting a systematic study of the dependence across iterations—and thereby understanding the mixing properties—is nontrivial. To sidestep the complexities of this dependence, one possible alternative is to explore uniform convergence. However, we have not pursued these directions rigorously and therefore prefer not to make any speculative claims.

Weakness-2, Question-3: KL divergence can be accommodated in Section 2.4.

We believe there is misunderstanding. KL divergence is a functional Bregman divergence, and so Theorem 2.4 of Section 2.4 covers KL divergence. See our discussion in Lines 212-214. Therefore, KL divergence is included in our main framework of FDE. However, when it comes to obtaining statistical error bound in Section 3, we need a separate theoretical analysis for KL divergence. We proved the convergence of our FDE based on KL divergence in Corollary B.9.

Weakness-3: Statistical tests and hyperparameter ablations

We have provided standard deviations (Atari games in Figures 1, 3-5) and error bars (LQR simulations in Figure 2) for the numerical results that allow assess for statistical significance. As for hyperparameter ablations, we indeed have varied the model complexity (the number of Gaussian mixture components as detailed in Appendix E.2.1.). For the experimental setups, we have included two levels of epsilon values for behavior policies (representing weak and strong satisfaction of Assumption C.3) and two different transitions (random and deterministic). Since we evaluated 7 Atari games, our experimental setup therefore includes $2$ (behavior policies) $\times$ $2$ (transitions) $\times$ $7$ (Atari games) $=28$ configurations. Each was tested with three levels of model complexity ($M=10,100,200$; see Lines 1490-1491), resulting in a substantial number of experiments. The outcomes are visualized in Figures 3-5. We believe this comprehensive evaluation sufficiently demonstrates our main conclusion that the methods developed within our framework are working.

Weakness-4: Proofs are all in Appendix.

Thank you for your suggestion. Based on your comment, we will add proof sketch/outline. Moreover, we will add a table of content in the supplementary material for easier navigation of the supplementary material.

References

[1] Peng et al. Statistical efficiency of distributional temporal difference learning. NeurlIPS. 2024.

[2] Wu et al. Distributional offline policy evaluation with predictive error guarantees. ICML. 2023.

[3] Duan et al. Risk bounds and rademacher complexity in batch reinforcement learning. ICML. 2021.

[4] Golowich et al. The role of inherent bellman error in offline reinforcement learning with linear function approximation. arXiv preprint. 2024.

Thank you very much for your efforts. Please find our responses to your comments below.

Weakness-1: Motivation; Value in light of FLE [1]

Thank you for your question. As emphasized in Section 2.1, a central goal of this work is to explore how to choose $\mathfrak{d}$ in order to develop a theoretically sound FDE method. Note that [1] focuses exclusively on the log-likelihood criterion (which is closely related to KL divergence), and it does not provide direct insight into this crucial question. To address it, we propose a theoretically motivated requirement for $\mathfrak{d}$ and show that functional Bregman divergences naturally satisfy this requirement, making them compelling candidates (see Theorem 2.4).

At a high level, different divergences respond differently to deviation between $\Upsilon$ and $\Psi_\pi$ (see Equation (5)), much like how various loss functions behave differently in empirical risk minimization. Therefore, the choice of divergence should be tailored to the specific problem at hand. Relying on a single divergence (such as KL) across all settings is unlikely to be optimal. In fact, there have been well-known criticism of KL divergence (e.g., unbounded divergence value, high sensitivity to the tails of distributions, necessity that one probability measure is absolutely continuous to the other, inability to consider closeness in outcome values), as discussed in [3], [4], [5]. Even in general statistical estimations, these limitations have motivated extensive research into alternative divergences beyond KL divergence [2].

To move toward a principled framework for choosing $\mathfrak{d}$ in specific applications, we believe the first step is to identify what constitutes a valid candidate. To date, this question remains largely unaddressed in the literature, and only a few valid choices have been identified - let alone systematically compared. Our work makes substantial progress on this front. While comparative analyses of certain functional Bregman divergences do exist and may provide practical guidance, we acknowledge that our paper does not yet offer a comprehensive decision-making framework for selecting among them in practice. We believe further investigation is needed to better understand the trade-offs between different functional Bregman divergences in the context of distributional OPE. Besides, identifying less common functional Bregman divergences that may be especially useful for distributional OPE is also interesting (e.g., Hyvarinen divergence which performed well in Atari games; see Tables 5-32). We believe they are important directions for future research. We will revise the paper to more clearly communicate these motivations and contributions.

Weakness-2: Many materials in the appendix.

Other reviewers have made some suggestions to improve the readability. We will revise the paper to address these comments. We plan to add more details (e.g., adding Theorem C.4 as suggested by Reviewer HoeH) to the main text. We also plan to add proof sketch/outline (as suggested by Reviewer vbgr). Besides, we will add a table of content in the supplementary material for easier navigation of the supplementary material.

Minor issues

Thank you very much for your suggestions and careful reading. Due to space limitation, we do not respond to them one by one. We will address these issues in the revised version.

Question-1: Examples of IPS metrics

In Section 3.1, the study is restricted to the case when the functional Bregman divergence is given as $\mathfrak{d}=m^2$ for some IPS metric $m$ (i.e., there exists an IPS metric such that the divergence is the squared IPS metric), and Table 1 presents examples of the functional Bregman divergence that can be written as a squared form of an IPS metric (Definition C.13). We have listed three examples of IPS metric in Appendix C.5, along with their corresponding inner products. Note that KL divergence is not a squared IPS metric and so is not included in Table 1, and the result of Section 3.1 (Theorem 3.3) does not apply to KL divergence. We have proved the convergence of FDE based on KL divergence separately in other results (Corollary B.9). (See Lines 317-318.)

Question-2: Relationship between $\eta$ and $\mathfrak{d}$

$\eta$ is a probability metric over $\mathcal{P}$, whose extension (either supremum-extension or expectation-extension of Equations 7 and 8) ensures that Bellman operator is a contraction (Definition 2.2). In our paper, we bound the inaccuracy of our estimation only based on the metric $\eta$ that satisfies Definition 2.2 (see Theorems 3.3, C.4, C.11), by using an $\mathfrak{d}$-based objective function (Equation 5). To successfully achieve this, we need a surrogate metric $m$ that links $\mathfrak{d}$ and $\eta$, which satisfies a series of properties (Assumption C.2 for Theorem C.4, Assumption C.9 for Theorem C.11) along with Assumption 3.2. Note that $m$ is used for quantifying the model complexity (see NS4 of Assumption C.2 and MS4 of Assumption C.9). One special case is using the divergence that has the form of $\mathfrak{d}=m^2$ with $m$ being an IPS metric (Definition C.13), but our results are not limited to squared IPS metics (as mentioned in Lines 310-311, 805).

We remark that the divergence in Theorem C.11 requires Assumption C.9, not Assumption C.2. We will fix this typo.

Question-3: Clarification of a proof step

Thank you for your careful reading. In the line before Equation 33, there was a typo. $\bar{m}_{n,1}$ should be replaced by its squared term. Then this solves the issue. As we mentioned in Line 1186, we can follow the logic of the proof for (14) to show the inequality. We will correct it in the revised version.

Before we procced, we would like to apologize for the typos listed above, which have caused unnecessary confusions. We will proofread the revised version carefully.

Question-4: Clarification of inaccuracy measure

We sampled 1000 state-action pairs from $d_\pi$ (Line 1509) according to Algorithm 1 of [6] (which we will clarify in the revised version). Since sampling from $d_\pi$ takes a long time for large $\gamma$, we saved the samples for each game (namely, pre-sampled observations) so that we can re-use them whenever we measure inaccuracy. To further clarify, we are not averaging Wasserstein-1 metric values for each pre-sampled observations of state-action pairs. Instead, we are calculating a single Wasserstein-1 metric value between the $S,A$-marginalized distributions in Line 1507 (which is also used by [1]). If necessary, we shall elaborate the algorithmic details of calculating this quantity.

Ancillary Question 1: Can we adopt $L_1$ form of expectation-extension?

Note that bounding $L_2$ version of expectation-extended metric (in the results of Theorems C.4, C.11) already implies bounding $L_1$ version of expectation-extended metric, by the fact that $L_1$ norm is bounded by $L_2$ norm for probability spaces.

But, as you mentioned, there is no need to exclude $L_1$ version of expectation-extension, which can also make the Bellman operator a contraction by the same proof structure based on Appendix C.3.1.

References

[1] Wu et al. Distributional offline policy evaluation with predictive error guarantees. ICML. 2023.

[2] Gneiting et al. Strictly proper scoring rules, prediction, and estimation. JASA. 2007.

[3] Korba et al. Statistical and Geometrical properties of the Kernel Kullback-Leibler divergence. NeurlIPS. 2024.

[4] Selten. Axiomatic characterization of the quadratic scoring rule. Experimental Economics 1.1. 1998.

[5] Bellemare et al. The Cramer distance as a solution to biased wasserstein gradients. arXiv preprint. 2017.

[6] Agarwal et al. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. JMLR. 2021.

Thank you very much for your efforts. Below are our responses to the comments you provided.

Weakness-1: Moving general results to the main text.

Thank you for your suggestion. Due to the page limit, we summarize our result in Equation 12, which accommodates other statements (Theorems C.4, C.11). Based on your comment, we plan to add more details (e.g., Theorem C.4) to the main text.

However, we would like to clarify two points. First, $p_{min}>0$ is not a strong assumption for the error bound guarantee in sup-norm metric (more specificially, $\eta_\infty$) in Theorem 3.3, since it is a strong uniform guarantee on every single state-action pair. For finitely many state action pairs (i.e., tabular setting in Theorem 3.3), $p_{min}>0$ is satisfied as long as every state-action pair has positive probability to be observed, as also assumed in [1], [2], [3]. Second, Theorem C.4 does not dominate Theorem 3.3. Although Theorem C.4 can accommodate all the functional Bregman divergences that can be applied in Theorem 3.3 (i.e., squared IPS metrics introduced in Section 3.1), they demonstrate slower convergence rates. This is explained in Lines 315-316. Indeed, Theorem 3.3 is significant as this shows near-optimal convergence rate for several divergences in tabular setting. Prior to our work, there is only one existing work [1] that shows optimal convergence rate (in tabular setting). See discussion after Theorem 3.3.

Weakness-2: Multi-modal returns

Note the return distributions depend on the choices of target policies. Moreover, we have also included experiments with modified environments with stochastic transitions (in addition to the original deterministic transitions), which obviously affect the shape of return distributions. Empirically, we observed examples that generate (marginal) return distributions that have multiple modes in our Atari game experiments.

Question-1: Finite horizontal setting

Our statistical error bounds (Theorems 3.3, C.4, C.11) do not apply to the finite horizontal settings. For finite horizontal setting, the algorithm will need to be modified. See, e.g., Section 2.1 of [4]. Note that there would be a set of return distributions defined over time, which are characterized by a series of distributional Bellman equations (which is different from that of the infinite horizon). And such relationship would be used for the finite horizontal extension, instead of the fixed-point characterization in the infinite horizontal setup.

Yet, we believe that a key part of work still provides insight for the construction of such distributional extension of the FQE. Based on a similar argument for (P2) and in Section 2.4, functional Bregman divergences are still expected to be natural candidates for finite-horizon extension.

References

[1] Peng et al. Statistical efficiency of distributional temporal difference learning. NeurlIPS. 2024.

[2] Rowland et al. An analysis of quantile temporal-difference learning. JMLR. 2024.

[3] Ma et al. Conservative offline distributional reinforcement learning. NeurlIPS. 2021.

[4] Wu et al. Distributional offline policy evaluation with predictive error guarantees. ICML. 2023.

Thank you very much for your efforts. Below are our responses to your comments.

Weakness-1: Too condensed.

We agree that we have a significant amount of content to fit into the paper. Other reviewers have made some suggestions to improve the readability. We will revise the paper to address these comments. We plan to add more details (e.g., adding Theorem C.4 as suggested by Reviewer HoeH) to the main text. We also plan to add proof sketch/outline (as suggested by Reviewer vbgr). Besides, we will add a table of content in the supplementary material for easier navigation of the supplementary material.

Weakness-2: Difficulty in interpreting the simulation results of Atari games

In Figure 1, except TVD (not a functional Bregman divergence) whose inaccuracy value stayed still throughout iterations, all other methods have at least shown improvement of accuracy. This demonstrates the validity of our FDE methods that are based on functional Bregman divergences. Furthermore, we can see that our FDE methods generally have lower mean inaccuracy values (dots in Figure 1) and more stable estimation (smaller standard deviation) than the baseline methods (FLE, QRDQN, IQN). Our settings can be categorized into four groups depending on random / deterministic transition and strong / weak coverage (explained in Lines 1542-1554). In the table below, we presented the mean of the rank values (1 being the best and 9 being the worst) of inaccuracies shown in Tables 5-32. We can see that our three FDE methods (KL, Energy, Hyvarinen) outperform baseline methods (FLE, QRDQN, IQN) in most settings.

Since the learned distribution is a conditional distribution, analyzing it requires examining many distributions—one for each state-action pair—which makes it difficult to compare them and identify patterns. That said, we appreciate your suggestions.

Weakness-3: Distinct policies in experiments.

In Appendix E.2.4, we have included four pairs of target and behavior policies for each of the seven Atari games. These include epsilon values being 0.3 (target) and 0.8 (behavior) for the deterministic transition (Lines 1542-1548), or 0 (target) and 0.5 (behavior) for the random transition (Lines 1549-1554). Although the target and behavior policies are both epsilon-greedy variants for the same DQN-trained policy $\pi_*$, the differences between these epsilon values amount to 0.5, which is large enough to make the target and behavior policies distinct. For each target policy, we have tried two choices of epsilon values for the behavior policy (slightly larger and much larger epsilon than the epsilon value for the target policy), to empirically verify the effects of coverage constant (Assumption C.3). Indeed, larger epsilon value for the behavior policy leads to larger inaccuracy levels, since they represent weaker coverage. This can be seen by comparing Row 1 and Row 2, or comparing Row 3 and Row 4 of Figures 3, 4, and 5.

Weakness-4: Total coverage of all state-action pairs.

Thank you for your comment. In (P2), we made a mistake in our writing-the uniqueness is only required up to the support of $\rho$ (i.e., almost everywhere w.r.t. $\rho$). We will revise the paper accordingly. Therefore, the updated (P2) does not require total coverage of all state-action pairs. Without total coverage, we can still achieve the convergence in expectation-extended metrics in Section 3.2, i.e., Theorems C.4 and C.11.

Question-1: Why Gaussian mixture model?

The Gaussian mixture model (GMM) is mainly chosen as a flexible model given its ability to approximate smooth densities well with sufficiently large number of components. Since a key prior work FLE [1] also used GMM in their simulations, we simply adopt the same choice in our experiments. For our experiments, certain divergences (among the variations of our FDE method) and existing method (like FLE [1]) require existence of density, which is conveniently satisfied by the use of GMM.

Question-2: Typos

Thank you for your careful reading. We will fix them.

References

[1] Wu et al. Distributional offline policy evaluation with predictive error guarantees. ICML. 2023.