Categorical Distributional Reinforcement Learning with Kullback-Leibler Divergence: Convergence and Asymptotics

We thank the reviewer for taking the time and effort in reviewing our paper, and we are pleased that they found our motivation compelling, appreciated our use of preconditioning to obtain a convergence guarantee, and found our use of asymptotic variance for the analysis to be a clear choice.

Comparison between Cramér and KL losses for learning distributions

We thank the reviewer for this suggestion, as we agree that an asymptotic comparison of the probability vectors learnt by the Cramér and KL losses is a valuable comparison to make. We report this result in a new proposition, which we state below. Let us define the iterates $p_k^{\text{Cram\’er-CMC}}$ as follows: $p_{k+1}^{\text{Cram\’er-CMC}} = p_k^{\text{Cram\’er-CMC}} + \alpha_k (h(G) - p_k^{\text{Cram\’er-CMC}})$. Then under the assumptions of Lemma D.3., we have that

$k^\beta ( p_{k}^{\text{Cram\’er-CMC}} - \tilde{p}^\pi) \overset{d}{\to} \mathcal{N}\left(0, \mathbb{E}_{G}[h(G) h(G)^T] - (\tilde{p}^\pi)(\tilde{p}^\pi)^T\right).$

We also want to emphasize that the goal of our work is not necessarily to argue that the KL version of the algorithm is always superior to the Cramer-based version (and indeed this is not always the case). Our main motivation is to study the KL version in its own right, which is the fundamental algorithm underlying the original deep RL implementation of distributional RL (C51; Bellemare et al., 2017).

Experimental comparisons between the KL loss and the Cramer loss under the function approximation setting

We appreciate that while our focus on the tabular setting potentially seems a limitation due to its simplicity, we want to emphasize that this was a conscious choice in order to develop a core understanding of categorical TD algorithms based on KL gradients, which have not been studied before, without confounding factors that often arise in larger-scale experiments (such as adaptive optimisation, function approximation, target networks, etc.).

We additionally note that the paper suggested by the reviewer performs an empirical comparison of KL-CTD and Cramér-CTD in the function approximation setting: C51 is the deep learning equivalent of KL-CTD, and S51 as used in their paper is the deep learning equivalent of Cramér-CTD.

Other Comments Or Suggestions

We thank the reviewer for bringing attention to these typos. We will be sure to fix them in the text.

References

Bellemare, M. G., Dabney, W., and Munos, R. A distributional perspective on reinforcement learning. In International Conference on Machine Learning, 2017.

We thank the reviewer for their time and effort in reviewing our paper, and for the helpful feedback and suggestions provided. We are pleased to hear that they found our work to provide valuable insights for practitioners and our idea of preconditioning to be concise and intuitive.

Asymptotic analysis missing for Cramér-CTD

We thank the reviewer for pointing out this limitation, and we have expanded the asymptotic results based on their feedback. In particular, we have derived additional results for the asymptotic normality of $p_k^{Cramer-CMC}$, and present this below.

Let us define the iterates $p_k^{\text{Cram\’er-CMC}}$ as follows: $p_{k+1}^{\text{Cram\’er-CMC}} = p_k^{\text{Cram\’er-CMC}} + \alpha_k (h(G) - p_k^{\text{Cram\’er-CMC}})$. Then under the assumptions of Lemma D.3., we have that

$$k^\beta ( p_{k}^{\text{Cram\’er-CMC}} - \tilde{p}^\pi) \overset{d}{\to} \mathcal{N}\left(0, \mathbb{E}_{G}[h(G) h(G)^T] - (\tilde{p}^\pi)(\tilde{p}^\pi)^T\right).$$

Extension of asymptotic analysis to general statistical functionals

We thank the reviewer for highlighting this as a direction to expand our current results, as we find this to be a nice generalization of our current analysis. We present the result for KL-CMC below, and will include it for TD and Cramér-CMC in the paper.

Suppose $\psi: \mathscr{P}(\mathbb{R}) \to \mathbb{R}^k$ is a statistical functional sketch, and let $J_\psi: \mathbb{R}^m \to \mathbb{R^k}$ be the Jacobian of $\phi \mapsto \psi(p^\phi)$. Then if $J_\psi$ is continuous in a neighbourhood of $\tilde{p}^\pi$, we have

$$k^\beta ( \psi(p_{k}^{\text{KL-CMC}}) - \psi(\tilde{p}^\pi)) \overset{d}{\to} \mathcal{N}\left(0, \frac12 J_\psi(\tilde{p}^\pi) \left(\mathrm{diag}(\tilde{p}^\pi) - (\tilde{p}^\pi)(\tilde{p}^\pi)^\top\right) J_\psi(\tilde{p}^\pi)^\top\right).$$

Work only considers the tabular case

We appreciate that while our focus on the tabular setting potentially seems a limitation due to its simplicity, we want to emphasize that this was a conscious choice in order to develop a core understanding of categorical TD algorithms based on KL gradients, which have not been studied before, without confounding factors that often arise in larger-scale experiments (such as adaptive optimisation, function approximation, target networks, etc.).

Missing blue line in Figure 2

The line for the MSE of TD is not missing, it rather perfectly overlaps with the line of Cramér-CTD, as predicted by the theory of Lyle et al. (2019). This is a key motivation for the work: while the behaviour of Cramer-CTD exactly coincides with classical TD learning, as predicted by Lyle et al. (2019), this shows that there are settings with KL-CTD has distinct, and sometimes superior, performance to these other algorithms, motivating us to understand the theory of this approach better. We will make this more emphasized in the text, and explore other methods of visualization to remove any confusion, such as making the linewidth of the blue curve larger.

References

Lyle, C., Bellemare, M. G., and Castro, P. S. A comparative analysis of expected and distributional reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2019.

We thank the reviewer for their time and effort they spent in reviewing our paper, and we aim to resolve their questions below.

Why do the authors consider the KL variant of categorical distributional RL, and what is the gap between it from the original algorithm, such as C51?

We focus on the KL variant of categorical distributional RL as it is exactly the fundamental algorithm used by C51. The gap between KL-CTD considered here and C51 is due to the use of policy evaluation instead of control, and tabular analysis instead of the deep learning setting.

The motivation for the study of KL-CTD is not clear in Section 3. Why do we not directly study the original algorithm?

We motivate KL-CTD in two ways in Section 3:

• Firstly, KL-CTD is the original categorical TD algorithm that C51 makes use of. This in itself is a strong motivation for its study.
• Our examples in Figure 2 and 3 serve as further motivation, showing that even in tabular settings, without neural network function approximation, KL-CTD can exhibit drastically different behaviour to earlier studied algorithms such as Cramer CTD and classical TD learning.

We emphasize that KL-CTD is the original CTD algorithm, as described at Eqn (6) and Line 141.

How to understand the convergence result in this paper based on KL divergence, given that distributional dynamic programming is only non-expansive if we directly employ the KL divergence to measure the distribution distance

Our convergence analysis relies on Lyapunov stability results, as opposed to writing our updates as a stochastic approximation of a contractive operator.

More specifically, there are several reasons why a result concerning the unprojected distributional Bellman operator's behaviour as measured by KL divergence is not pertinent to our analysis here: We are concerned with convergence of an algorithm that maintains a categorical representation of the return distribution, but the result above is concerned with the unprojected distributional Bellman operator. While we analyze an algorithm that is defined via the KL loss, this does not obligate us to make use of contraction results measuring distance in KL. In Bellemare et al. (2023, Chapter 5), this difference is emphasised by distinguishing between metrics used for analysis, and metrics used to define algorithms. In our case, we study an algorithm that uses KL between categorical distributions in its definition, and our analysis makes use of both KL (to measure distance from the fixed point), and Cramér distance (we make use of contractivity of the categorical distributional Bellman operator in Cramér distance to show that the KL is a Lyapunov function; see the proof of Proposition 5.4 and line 783 in particular).

Please let us know if you have any further queries on this point.

$\tilde{p}$ above Equation (8)

The $\tilde{p}^\pi_i$ appearing in the equation $\tilde{\eta}^\pi = \sum_i \tilde{p}^\pi_i \delta_{z_i}$ represents the probabilities of the fixed point of the projected distributional Bellman operator.

​​ It is not clear why C is introduced in Eq.9 and 10. Could any intuitive explanation be provided?

Our previous calculation on lines 225-227 shows that the KL divergence from the current estimate to the fixed point may be increasing under the KL-CTD dynamics. But from this expression, in light of Proposition 5.3., we can reverse-engineer a modification to the updates which would guarantee the decrease of the KL, and lead to convergence. This modification is exactly the preconditioning using $C^\top C$.

It seems that the superiority of the proposed algorithms PKL-CTD over KL-CTD and TD highly depends on the step size.

The reviewer is correct that which algorithm is best between the various considered in Figure 7 depends on the step size, however we would like to emphasize that at least in the settings in the figure, the dependence of PKL-CTD on step size has a generally equal or smoother curve than the dependence of KL-CTD and TD on step size, which both have deep learning equivalents which are highly used in practice (DQN and C51).

Why consider the asymptotic normality/variance in value estimates in Propositions 6.1 and 6.4? What is the main regularity condition? Is there any potential application of these asymptotic properties?

We study the asymptotic distributions of the value estimates because they are able to give us exact forms of the errors incurred by the algorithms. These asymptotic results have led to the bias/variance results in our paper, and we believe that they have further applicability for deriving insights in future work.

Comparison with Rowland (2018)

While Rowland et al. (2018) study categorical dynamic programming and the Cramer-CTD update, this paper focuses on the categorical TD learning algorithm based on the gradient of the KL loss, matching the form of loss used in the original categorical distributional RL paper (Bellemare et al., 2017).

We thank the reviewer for their time and effort in reviewing our paper, and for the helpful feedback provided. We are pleased that the reviewer found our work on the convergence of PKL-CTD, the asymptotic variance analysis, and the connection between theory and empirical observations to be valuable. Below, we address each of the reviewer’s points in detail.

Analysis is restricted to synchronous policy evaluation scenarios

Synchronous/Asynchronous

We note that all experimental results in Section 7 are done using asynchronous updates (with synchronous updates performed as an ablation in Appendix F), though the reviewer is correct that our convergence result addresses the synchronous setting. In light of the reviewer’s suggestions we have expanded our analysis, and give a sketch below of how it can be applied to analyse asynchronous cases too.

Following work presenting core convergence guarantees in asynchronous stochastic approximation, such as Borkar (2008), considering the case where states are updated according to an ergodic Markov chain with stationary distribution $c$, the key concern to establishing convergence under appropriate technical conditions on step sizes is to establish a corresponding Lyapunov function for the continuous-time dynamics

$\partial_t \phi_t (x) = c(x) C^\top C (T^\pi - I) p^{\phi_t}(x),$

which take into account the average per-state weighting of updates. From our analysis of the synchronous case (see Proposition 5.4), it can then be verified that $L(\phi) = \sum_{x\in\mathcal{X}} \frac{d^\pi(x)}{c(x)} \mathrm{KL}(\tilde{p}^\pi(x) \| p^\phi(x))$ is a Lyapunov equation for this ODE, which can then be used to guarantee its convergence.

Policy Evaluation/Control

We focus on the policy evaluation setting in this paper, and leave the control setting as interesting future work. We remark that understanding the policy evaluation setting first is typical in the analysis of both standard and distributional reinforcement learning.

Proof techniques used are relatively standard, existing stochastic approximation theory

We agree that our proof of Theorem 5.5 uses existing stochastic approximation theory, although we would argue that many applications of stochastic approximation theory to RL make use of core existing theory (such as in the classic papers analyzing tabular TD algorithms by Jaakkola et al. (1994) and the linear TD analysis of Tsitsiklis and Van Roy (1997)), and indeed novel theory and techniques of stochastic approximation are often published as stochastic approximation results in their own right.

Further, we want to emphasize that our theoretical analysis leads to a number of novel results, such as an undiscovered algorithm (PKL-CTD), the effective learning rate phenomenon which we show is fundamental to KL-based categorical algorithms, and the bias-variance tradeoff in the atom locations.

Convergence of non-preconditioned KL-CTD

We agree that the reviewer is entirely correct that our counterexample to the KL being a Lyapunov function for non-preconditioned KL-CTD does not exclude its possibility of convergence, and we will make this clearer in the paper. We chose to include this as an additional result in the appendix as it is a natural question in light of Proposition 5.4 (does the weighted KL, our exhibited Lyapunov function for PKL-CTD, also work as a Lyapunov function for KL-CTD?), we did not intend to suggest that non-preconditioned KL-CTD cannot converge.

Recent works on quantile-based DRL are not cited

We would like to ensure that no relevant work is missing from our discussion of related work in Appendix E. We invite the reviewer to suggest any works in particular they find to be missing.

Other Comments Or Suggestions

We thank the reviewer for all of the additional catches/suggestions, and we will ensure to fix them in the text.

References

Jaakkola, Tommi, Michael Jordan, and Satinder Singh. "Convergence of stochastic iterative dynamic programming algorithms." Neural Computation (1994).

Tsitsiklis, John and Van Roy, Ben "An analysis of temporal-difference learning with function approximation," in IEEE Transactions on Automatic Control (1997)

Borkar, Vivek. Stochastic approximation: a dynamical systems viewpoint. Cambridge University Press, (2008).