A Distributional Analogue to the Successor Representation

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.

The significance of section 5.2 is unclear…

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.

Thank you for your thorough review and the time you have dedicated to our paper. We have carefully addressed each of your questions and concerns in our rebuttal. We hope our clarifications align with your expectations, and we are eager to answer any further questions you have.

Kind regards, Authors.

We thank the reviewer for their kind comments and overall positive review. We agree with the reviewer that the DSM is novel and that insights from our $\delta$-model method are useful beyond estimating the DSM. We hope that our response below helps provide additional motivation for the problem we introduce and inspire exciting future work on the DSM. We look forward to an engaging discussion period.

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. 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.

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.

Thank you for your thorough review and the time you have dedicated to our paper. We have carefully addressed each of your questions and concerns in our rebuttal. We hope our clarifications align with your expectations, and we are eager to answer any further questions you have.

Kind regards, Authors.

In response to the reviewer’s concern that our method may be limited to relatively small discount factors, we launched a $\delta$-model training run in the Pendulum environment with a longer horizon ($\gamma = 0.99$). We used a n-step bootstrapping estimator for the $\delta$-model targets with $n=15$. We found that, even in the early stages of training, the return distribution predictions from the DSM are quite accurate, and have very similar quality to those reported for $\gamma = 0.95$. You may find the corresponding plot in the supplementary material, under the name dsm-gamma-099.pdf. We hope this alleviates the reviewers’ concerns about the feasibility of training $\delta$-models for longer horizons.

We thank the reviewer for their time and positive feedback. We agree with the reviewer that the DSM is a promising and important direction of research. We address the questions outlined by the reviewer below and would like to reiterate that traditional distributional RL and the DSM are solving different problems and aren’t directly comparable. We look forward to engaging with the reviewer during the rebuttal period and hope the reviewer may consider revising their score based on our discourse.

For distributional RL papers, a common benchmark is visual input environments …

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.

As for Eq.7, are there any requirements for the reward function …

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.

Can you further discuss the relationship between previous work like Zhang et al 2021b …

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.

Is r assumed to be known in experiments?

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.

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?

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.

References

• Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer. "High-Resolution Image Synthesis with Latent Diffusion Models." IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
• Patrick Esser, Robin Rombach, and Björn Ommer. "Taming transformers for high-resolution image synthesis." IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
• Danijar Hafner, Timothy P. Lillicrap, Mohammad Norouzi, Jimmy Ba. "Mastering Atari with Discrete World Models." International Conference on Learning Representations (ICLR), 2021.
• Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy P. Lillicrap, David Silver. "Mastering Atari, Go, chess and shogi by planning with a learned model." Nature Vol. 588, Pages 604-609, 2020.

Thank you for your thorough review and the time you have dedicated to our paper. We have carefully addressed each of your questions and concerns in our rebuttal. We hope our clarifications align with your expectations, and we are eager to answer any further questions you have.

Kind regards, Authors.

We thank the reviewer for their insightful comments and engagement with our work, and are glad to hear that the reviewer finds the DSM novel, and the $\delta$-model an innovative means of learning the DSM.

We view the reviewer's main concerns as relating to the possibility of theoretical analysis of $\delta$-models, and the possibility of investigating alternative approximations to the DSM. We believe that both of these are important directions for future work. We also believe that our introduction of the DSM learning problem, together with the empirically successful approach of training $\delta$-models, are already substantial contributions in their own right. We respond to all comments in detail below, and eagerly await further discussion on these points.

Firstly, it remains unclear why equal weights are utilized in equation (10)...

Our choice to use equally weighted particles to approximately represent probability distributions is motivated by several considerations:

• Prior work in distributional reinforcement learning. Both the quantile representation introduced by Dabney et al. (2018), and the equally weighted particle representation introduced by Nguyen-Tang et al. (2021) have this form.
• Prior work in computational statistics. Using equally weighted particles to approximate distributions is also a common approach in computational statistics, such as Wasserstein/MMD gradient flow approximations (see e.g. Chizat and Bach, 2018, Arbel et al., 2019) and herding (Chen et al., 2010).
• This parameterization can still effectively assign weight to regions of the space of occupancy measures. For instance, if a given measurable set of the space of occupancy measures should have more probability mass, the trained $\delta$-model can (and does) place more model atoms in that measurable set.

We agree with the reviewer that this is not the only choice of approximate representation that could be made. We view this work, introducing the distributional successor measure, as opening up a new research question within distributional RL: how to approximate and learn the DSM effectively. The delta-models introduced in this paper, together with the experimental results, provide an initial answer to this question, and we expect future work to continue to develop understanding in this area.

Secondly, there is a lack of guidance on selecting the value of m in equation (10)...

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.

Let's say we have $X = 0,1$… [are there] any theoretical guarantee[s]?

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.

Thirdly, there appears to be a dearth of theoretical justification …

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.

Thank you for your thorough review and the time you have dedicated to our paper. We have carefully addressed each of your questions and concerns in our rebuttal. We hope our clarifications align with your expectations, and we are eager to answer any further questions you have.

Kind regards, Authors.