Paper: A Distributional Analogue to the Successor Representation (ICLR-2024)

Reviewer comment (Theory):
A few concepts are quite confusing along the paper. First, why is it necessary to redefine the occupancy measure/state(-action) distribution as successor representation or successor measure?

Author response:
We thank the reviewer for their constructive feedback. We suspect that there has been a fundamental misunderstanding about what the DSM represents, and we hope to clarify this here. Particularly, we would like to address the following statements from the review:

> The successor measure is defined as, conditioned on an initial state, the Dirac distribution conditioned on a random trajectory…

> ...why is it necessary to define the new "random occupancy measure", which is a distribution over distribution, but the inner distribution is just the Dirac distribution which is determined by the realization of the sample of the outer distribution?

The DSM is not a Dirac distribution conditioned on a random trajectory, and the “inner distributions” are also not Dirac distributions. The DSM is a probability distribution over distributions over the state space. Notably, the DSM has support on many distributions over the state space, and the distributions on its support are supported generally on the whole state space. Perhaps the reviewer is referring to the $\delta$-model parameterization of the DSM; however, the claims are still not accurate:
- $\delta$-models are supported on a collection of Diracs, not just one. That is, the $\delta$-models can be supported on arbitrarily many occupancy measures (controlled by the number of model atoms). It is important to note that while these are Diracs, they are Diracs each located on a probability distribution over the state space. So, the locations of the model atom Diracs are infinite dimensional.
- Each model atom is represented as a generative model of state. Our architecture permits iid sampling from these model atoms, which enables us to leverage the unbiased sample-based MMD estimator to learn the locations of the model atoms.
- Thus, there are two layers of samples. Sampling from a $\delta$-model produces a generative model (one of the particles of the $\delta$-model), and the generative model samples themselves can draw arbitrarily-many iid state samples to represent the corresponding occupancy measure.

> ...why is it necessary to redefine the occupancy measure/state(-action) distribution…

As the reviewer points out, these concepts are tightly related to the occupancy measure. However, this is standard nomenclature in the literature; the successor representation/measure is often referred to as the discounted occupancy measure.

Reviewer comment (Theory):
Why is it necessary to define the new "random occupancy measure", which is a distribution over distribution, but the inner distribution is just the Dirac distribution which is determined by the realization of the sample of the outer distribution? From my understanding, if one wants a distribution over the value function, suppose they already have the occupancy measure, they could easily define the distribution of the value function (which is in $\Delta(\mathbb{R}^1)$) by projecting the occupancy measure to the reward vector. I think this is also indicated by prop. 1: to define the distributional return, one needs to take one expectation over the random occupancy measure.

Author response:
We thank the reviewer for their constructive feedback. We suspect that there has been a fundamental misunderstanding about what the DSM represents, and we hope to clarify this here. Particularly, we would like to address the following statements from the review:

> The successor measure is defined as, conditioned on an initial state, the Dirac distribution conditioned on a random trajectory…

> ...why is it necessary to define the new "random occupancy measure", which is a distribution over distribution, but the inner distribution is just the Dirac distribution which is determined by the realization of the sample of the outer distribution?

The DSM is not a Dirac distribution conditioned on a random trajectory, and the “inner distributions” are also not Dirac distributions. The DSM is a probability distribution over distributions over the state space. Notably, the DSM has support on many distributions over the state space, and the distributions on its support are supported generally on the whole state space. Perhaps the reviewer is referring to the $\delta$-model parameterization of the DSM; however, the claims are still not accurate:
- $\delta$-models are supported on a collection of Diracs, not just one. That is, the $\delta$-models can be supported on arbitrarily many occupancy measures (controlled by the number of model atoms). It is important to note that while these are Diracs, they are Diracs each located on a probability distribution over the state space. So, the locations of the model atom Diracs are infinite dimensional.
- Each model atom is represented as a generative model of state. Our architecture permits iid sampling from these model atoms, which enables us to leverage the unbiased sample-based MMD estimator to learn the locations of the model atoms.
- Thus, there are two layers of samples. Sampling from a $\delta$-model produces a generative model (one of the particles of the $\delta$-model), and the generative model samples themselves can draw arbitrarily-many iid state samples to represent the corresponding occupancy measure.

> ...why is it necessary to redefine the occupancy measure/state(-action) distribution…

As the reviewer points out, these concepts are tightly related to the occupancy measure. However, this is standard nomenclature in the literature; the successor representation/measure is often referred to as the discounted occupancy measure.

> From my understanding, if one wants a distribution over the value function, suppose they have already have the occupancy measure, they could easily define the distribution of the value function (which is in P(R)) by projecting the occupancy measure to the reward vector.

> ... First, the paper should compare with other distribution RL methods, …

Our goal is far more ambitious than learning a distribution over the return. Indeed, this can be accomplished with standard distributional RL techniques like you suggest. Moreover, with just the occupancy measure, “projecting onto the reward” only produces a point estimate of the return, not a distribution – indeed, this is how transfer with the SR works, and this process is shown in equation (4) of the paper. Our approach is modeling another “axis of distribution” beyond the SR/occupancy measure, and that is how we can “project” our DSM onto a reward function to produce a distribution over returns. To achieve this functionality, it is necessary to properly model the interaction between random occupancy measures, and no other method has done this.

Reviewer comment (Presentation):
The significance of section 5.2 is unclear. How to tune MMD does not directly relate to the significance of the paper, and many descriptions seem not significant (for example, the detailed description of the median trick).

Author response:
As for the significance of Section 5.2, this is a good question. What distinguishes our problem setting from those where MMD optimization is commonly applied is that, due to bootstrapping, the target distributions that we model are non-stationary, and they can change over time (especially towards the beginning of training). Even with static target distributions, kernel optimization tricks are commonplace as you pointed out. However, existing RL algorithms based on MMD optimization (e.g., Nguyen-Tang et al. (2021), Zhang et al. (2021)) did not address this problem. Especially given the complex space of distributions being modeled in our work, we hypothesized that careful treatment of adaptive kernel parameters could substantially improve the quality of the resulting $\delta$-models. We tested this hypothesis by ablating the adaptive kernel, and found that this was a crucial factor in the success of $\delta$-model training – we invite the reviewer to consult Appendix D (particularly, Figure 9, right), for these results.

### References
- Thanh Nguyen-Tang, Sunil Gupta, Svetha Venkatesh. "Distributional reinforcement learning with maximum mean discrepancy." AAAI Conference on Artificial Intelligence (2021).
- Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, Tie-Yan Liu. "Distributional reinforcement learning for multi-dimensional reward functions." Neural Information Processing Systems (NeurIPS), 2021.

Reviewer comment (Experiments):
The experiment benchmark (windy grid world) is somehow toy compared with other distributional RL papers.

Author response:
Regarding learning a DSM from pixels, one of the reasons we modeled lower-dimensional state spaces is that we believe these results are indicative of how the method will perform if one were to scale to high-dimensional observations for the following reason: We believe learning high-dimensional data distributions isn’t the best way to scale a method like DSM. Compounding the difficulty of learning a distribution over images, DSM must model a distribution over distributions of images. As evidence of this claim, performing dimensionality reduction and operating in a structured low-dimensional latent space seems to yield state-of-the-art results in model-based RL (e.g., Schrittwieser et al. 2020, Hafner et al. 2021) and generative modeling of images (e.g., Esser et al. 2021, Rombach et al. 2022). Given such latent embeddings, our current $\delta$-models can learn the DSM as shown in our paper. Thus, we believe that supplementing DSM with a dimensionality-reduction component is out of the scope of this paper, but would make for a very exciting follow-up work.

Reviewer comment (Theory):
As for Eq.7, are there any requirements for the reward function $r$ besides deterministic? For instance, does it require $r$ to be linear? BTW, is $r$ assumed to be known in experiments?

Author response:
Regarding the requirements of the reward function, this is a strength of the DSM model: all that we require is that the reward function is measurable, which encompasses essentially any reward function. In particular, linearity is not required. The reward function can be a black box; all we need is the ability to query the reward. 

In our experiments, we assumed access to reward function queries (as is commonplace with SR models). There are several instances where this can be applied, such as goal-reaching,  target-tracking tasks, or continuous control tasks that require engineered reward functions. However, more generally, we can also learn a reward function from data, for instance, by supervised learning from reward samples, using IRL from expert trajectories, or learning a reward function from preferences as in RLHF.

Reviewer comment (Novelty):
Can you further discuss the relationship between previous work like Zhang et al 2021b, which also learn multi-dimensional return distribution via MMD?

Author response:
The work on distributional RL with multi-dimensional reward functions (e.g., Zhang et. al 2021b) is indeed closely related, as the reviewer suggests. Multi-dimensional distributional RL methods can simply model the distribution over random finite-dimensional vectors corresponding to returns from reward signals seen at train time. As such, one could use a multi-dimensional return distribution function like that learned by Zhang et al. to infer the return distributions for linear combinations of the reward functions that are used for training. On the other hand, the DSM is modeling a complex distribution over the simplex on the state space, which requires a novel approach to distribution estimation using generative models, as highlighted by Reviewer 9gwA. Notably, it is this distinction that allows the DSM to predict return distributions for arbitrary bounded measurable reward functions.

Reviewer comment (Experiments):
For distributional RL papers, a common benchmark is visual input environments like Atari. I think current benchmark (windy gridworld) is somehow toy. It will be helpful to see experiments with larger scale. What’s more, besides zero-shot policy evaluation, another advantage of SR is multitask training. Can the proposed method be combined with the multitask training?

Author response:
Regarding learning a DSM from pixels, one of the reasons we modeled lower-dimensional state spaces is that we believe these results are indicative of how the method will perform if one were to scale to high-dimensional observations for the following reason: We believe learning high-dimensional data distributions isn’t the best way to scale a method like DSM. Compounding the difficulty of learning a distribution over images, DSM must model a distribution over distributions of images. As evidence of this claim, performing dimensionality reduction and operating in a structured low-dimensional latent space seems to yield state-of-the-art results in model-based RL (e.g., Schrittwieser et al. 2020, Hafner et al. 2021) and generative modeling of images (e.g., Esser et al. 2021, Rombach et al. 2022). Given such latent embeddings, our current $\delta$-models can learn the DSM as shown in our paper. Thus, we believe that supplementing DSM with a dimensionality-reduction component is out of the scope of this paper, but would make for a very exciting follow-up work.

With regard to multitask training, can you please clarify or give an example of the type of setting you had in mind? Like SR methods, our DSM algorithm is policy-dependent but is otherwise task-agnostic. Our response to Reviewer 9gwA also discusses this, but please let us know if this does not address your question.

Reviewer comment (Theory):
Similar to the conventional successor feature and successor measure, the decoupling of the transition structure and reward functions still assumes a fixed policy and transition dynamic, limiting the usefulness of the proposed method.

Author response:
The reviewer makes a fair claim that the theory of DSM (and other SM/SR tasks) is limited to fixed policies. Despite this, the successor representation has garnered significant interest in many facets of RL, for example, representation learning, exploration, and temporal abstractions which could conceivably benefit from a distributional perspective; examples are given in the second paragraph of Section 7.

Our principal goal with this work was to address the difficulty of representing and learning a DSM, which has not been done previously, and we suspect that future work can devise control methods based on the DSM. There are some significant challenges in this direction, which are not specific to the DSM model itself:
- Even in distributional RL, there is generally no guarantee that return distributions will converge under greedy distributional updates.
- In risk-sensitive optimal control, it is generally necessary to employ policies that are either non-stationary or non-Markovian, and it is not immediately clear how to model occupancy measures or perform zero-shot evaluation for these types of policies.
- The work of Touati et. al (2021) introduces the ‘forward-backward’ representation of the successor measure, which enables an agent to learn the SM for all optimal policies simultaneously. But in the distributional case, this necessitates a definition of ‘optimal policy’, which has its own challenges, as mentioned in the previous point.

Solving any of these problems would be an incredible contribution in their own right, and we leave this to future work.

Reviewer comment (Novelty):
The usefulness of the distributional SM is quite limited at this point. I would recommend discussing more about the potential applications of the learned distributional SM other than the zero-shot distributional policy evaluation.

Author response:
That said, we believe there are exciting applications for DSM as presented in this work. Unlike any other algorithm we are aware of, the DSM allows us to generalize across tasks and risk-sensitive criteria which has many potential applications; one such example is the zero-shot risk-sensitive policy selection procedure presented in Section 6.

Reviewer comment (Experiments):
Will the setting of $\gamma = 0.95$ limit the usefulness of the proposed method in practice when we care about the return of a long episode?

Author response:
> Will the setting of $\gamma=0.95$ limit the usefulness of the proposed method in practice when we care about the return of a long episode?

Regarding the limitation of the setting of the discount factor in our experiments, we are happy to report results for larger discount factors – we are running these experiments and will report back. Note, however, that the DSM already handles larger horizons more gracefully than the comparable $\gamma$-models: in their experiments, they could not achieve reasonable performance by training with $\gamma=0.95$ directly; instead, they found it necessary to introduce a heuristic schedule for $\gamma$. Beyond that, even with the discount schedule, they could only achieve good results with a density model, which required further assumptions on the MDP dynamics. On the other hand, thanks to the n-step bootstrapping approach to generative TD learning that we introduce in this work, we can successfully learn $\gamma$- and $\delta$-models for $\gamma=0.95$ (and we will confirm the same for $\gamma=0.99$) without any heuristics like the discount schedule or additional assumptions about the dynamics.

Reviewer comment (Experiments):
Secondly, there is a lack of guidance on selecting the value of  $m$ in equation (10), or determining how many $m$ values are required.

Author response:
Trained $\delta$-models better approximate the true DSM as $m$ increases. Figure 9 in Appendix C shows that the Wasserstein distance between the DSM and the MC return distribution indeed decreases as we increase the number of model atoms.

Reviewer comment (Theory):
Thirdly, there appears to be a dearth of theoretical justification for the effectiveness of this modeling and approximation approach. While equation (10) may seem suitable if it exactly represents the true SDM, practical implementation would not align with this ideal scenario.

Author response:
The reviewer is correct that convergence theory for the algorithm has not been established. Our approach is based closely on approximate dynamic programming methods, as well as an MMD-based loss whose behavior in the case of fixed targets has been rigorously analyzed (Arbel et al., 2019). In particular, in our work, we have established that the distributional SM can be computed by distributional dynamic programming (see Appendix F.1), as is the case of the return distribution function in standard distributional RL. What is missing is the behavior of this algorithm under stochastic approximation of the distributional Bellman operator, as well as the consequences of projecting the targets onto the space of EWP representations.

The quantile TD-learning framework and QR-DQN algorithm of Dabney et al. (2018), which is closely related to our approach, had also not established this theory, which was only very recently resolved by Rowland et al. (2023). This has been a recurring trend in distributional RL, where new important algorithms and their analyses have often come separately. Notably, distributional RL with EWP representation and an MMD objective was employed by Nguyen-Tang et al. (2021) and the multidimensional distributional RL work of Zhang et al. (2021), despite the (ongoing) lack of convergence results in both works. We are definitely in agreement with the reviewer that theoretical analysis of these algorithms is a crucial direction for distributional RL as a whole, and we strongly believe that our contribution presents substantial value to the community in providing a starting point for future theoretical and practical innovations, as earlier works in distributional RL have done in the past.

Reviewer comment (Theory):
Let's say we have X = {0,1} (binary). Then, a set of $P(X)$ is parametrized by just one parameter $\mu \in [0,1]$. So, $P(P(X))$ is a set of distributions over $[0,1]$. So, if I understand correctly, learning SDM is equivalent to estimating a distribution over  $[0,1]$. Even in this simple case, does the author's approach have any theoretical guarantee? (finite $m$ and equal weights look restrictive?)

Author response:
The reviewer’s understanding of the DSM here is correct. It is true that we cannot necessarily model the exact DSM with the EWP parameterization – indeed, it is impossible to guarantee this, since the space of DSMs is far too large to search and represent. Our goal here, like in distributional RL, is to find a good approximation to the DSM among a tractable model class (which is the space of EWP representations here). Further discussion about convergence is given below.

Reviewer comment (Novelty):
Equation (4) may appear somewhat elementary to researchers within the RL community. The current phrasing, such as "Blier et al. 2021 derived this equation...," seems inaccurate and should be revised. I believe that this equation had already gained widespread recognition prior to the work of Blier et al. (2021), as it is commonly featured in numerous standard RL texts and papers.

Author response:
We agree with the reviewer that the discrete case, i.e., $V^\pi = (1-\gamma)^{-1}\Psi^\pi r^\pi$ has been broadly recognized, but the continuous measure-theoretic case hasn’t been as widely discussed. We looked into the origins of this equation in the measure-theoretic case, and it seems that both Blier et al. 2021 (Equation 11) and Janner et al. 2020 (Equation 6) had independently introduced this equation. We have revised the text to reflect this. If the reviewer is aware of any prior work beyond these papers, we would greatly appreciate this and will revise the text accordingly.

Reviewer comment (Novelty):
As a related point, I generally believe the author should not emphasize whether the definition pertains to discrete or continuous spaces to such an extent. The transformation from a discrete space (when a base measure is the counting measure) to a continuous space (when the base measure is a Lebesgue measure) is typically straightforward for individuals with a basic understanding of probability. Therefore, in Section 3.2, the statement "though this result is novel in the case of more general state spaces" may be somewhat misleading. I suggest that this aspect should not be categorized as "novel." Instead, I recommend that the author simply highlight the distinctions between the two contexts (SDM and standard distribution RL with rewards).

Author response:
This definition, as well as the following definitions, are given in the general measurable state space setting (it applies to both discrete and Polish state spaces). As you suggest, this definition when invoked with a discrete state space endowed with the counting measure reduces to the familiar definition of the SR and related identities. We agree that this can and should be made more clear, and we have adjusted the paper to address this.

With regard to the claim of novelty, we agree that perhaps we did not appropriately emphasize how the proposition is novel. As you claim, generalizing notions on discrete spaces to more general measurable spaces (particularly Polish spaces) is a fairly standard exercise. The novelty is in the formulation of the DSM as the fixed point of a distributional recurrence relation. While distributional Bellman equations for the (scalar) return have been known since at least as far back as Sobel (1982), a formulation of reward-agnostic distributional dynamic programming for the distributional successor measure has not previously been proposed. Upon formulating something like Equation 8, we agree that mathematically its correctness is not overly difficult to prove. As you suggest, it is the distinction between this object and the distribution over multidimensional returns which is new, and we have adjusted the text to reflect this.  We greatly appreciate the reviewer’s attention to detail here.

Reviewer comment (Reproducibility):
It is somewhat unclear which parameters are precisely optimized throughout the entire algorithm. As I understand it, the author optimizes parameters for all generative models simultaneously in equation (16). It would be beneficial to present the algorithm using an algorithmic environment for clarity.

Author response:
Yes, you are correct that we are optimizing the parameters of all the model atoms via Equation 16. Furthermore, Algorithm 1 in Appendix 1 details the complete algorithm which we hope can clear up any confusion.