Proto Successor Measure: Representing the space of all possible solutions of Reinforcement Learning

We thank the reviewer for their comments on the paper. We are encouraged to see that the reviewer found the idea interesting, well-motivated, and theoretically sound. We address the weaknesses and issues pointed out by the reviewer below:

This paper looks like it has been created in haste, with nonexistent figure references in line 257 and 279, a relatively large whitespace left on page 9 after the conclusion, and a missing reproducibility statement.

We apologize for missing figure references and have corrected them now. The code to reproduce our results can now also be found in supplementary material. We have revised the paper to be more comprehensive based on the reviewer’s feedback.

Evaluation and Ablations

First of all, only a maze environment and Fetchreach are used, which does not correlate to "a number of environments".

We have added 16 continuous control tasks from DM-control for evaluating and comparing PSM to baselines. This complements our discrete domain experiments and establishes that PSM indeed outperforms prior methods on a number of environments and tasks.

Adding to this, in Section 5.2 of [Ahmed Touati & Yann Ollivier, 2021], it seems their Forward-Backward (FB) algorithm easily solves the Fetchreach environment. Can the authors comment on this performance discrepancy? Also, for completeness and to say you outperform prior baselines, the authors could at least add the Ms. Pacman environment as done in prior work.

The Forward-backward algorithm from [Ahmed Touati & Yann Ollivier, 2021] evaluates grid worlds, discretized Fetchreach, and Pacman. Unfortunately, their datasets were never publicly released, contributing to difficulties in reproducing their setup exactly. Our results are based on the datasets we collected by running the exploration algorithm RND on fetchreach environments and enumerating all transitions for discrete maze and gridworlds. We have now added evaluations on 16 more tasks on 4 continuous challenging environments.

There is no real discussion or future work component in this paper, although there is space for it.

We thank the reviewer for their comment; we have added a discussion on limitations and future work in the conclusion section of the paper.

Fig.1: Figure 1 should be trimmed. The part of the image with "Find a Behaviour in a Dynamical System" is not necessary and confusing.

Our intent is to demonstrate the zero-shot nature of RL, that arises due to the representations learned by PSM. i.e. We aim to demonstrate that given any task requested by the user, we can obtain an optimal policy just by searching in the affine polytope represented by PSM. We have made changes to the figure to make it more clear.

References: The references are organized very poorly

We appreciate the reviewers' advice and have updated all the references to point to the published venue as opposed to their arXiv citations.

Why are there no ablations, and no other quantitative evaluations other than the Fetchreach environment?

We have now added evaluations on a larger set of more complicated tasks in the paper. We are running experiments ablating the representation dimensionality of PSM which will be added to the paper before the discussion period ends.

While we have striven to make the evaluations more comprehensive by adding more environments and tasks, we believe that our theoretical contributions might have been overlooked. We hope that the reviewer can reassess our paper in the light of a more thorough evaluation and writing improvements. Please let us know if you have any more questions.

We thank the reviewer for raising their score.

1. We respectfully disagree that our paper was more similar to an abstract by the paper submission deadline. Prior published works [1,2] in similar areas that focused on theoretical insights often demonstrated their empirical results in mostly discrete environments like ours. We believe our paper consisted of impactful and novel theoretical insights. We clarify that our initial experiments were on 3 environments (~30 tasks, tasks are defined by reward function or goals), and we have subsequently added 4 environments (16 tasks). Since our method is zero-shot RL, a number of tasks are evaluations of a variety of reward functions for every environment. We also clarify that we have added one ablation since the initial draft. In our opinion, the additional experiments should not say the initial experiments were lacking but shows our commitment to showing the method is indeed empirically promising on different environment set.
2. We agree that we missed two references in the final draft that arose on our part due to last moment latex compilation error. Citing arXiv papers instead of published version of the paper has been adopted by many papers published recently (eg. Direct Preference Optimization (NeurIPS 2023 best paper): 11, 18, 46 references point to arXiv despite being old published papers). We still believe in citing the published versions and had updated all our references. We believe it is a bit harsh to penalize us for this since a number of published papers have grammatical, vocabulary, and arXiv citation errors, and these aspects can be improved with reviewers' comments.

[1]: Diego Gomez, Michael Bowling, Marlos C. Machado; Proper Laplacian Representation Learning, ICLR 2024.
[2]: Lucas N. Alegre, Ana L. C. Bazzan, Bruno C. da Silva; Optimistic Linear Support and Successor Features as a Basis for Optimal Policy Transfer, ICML 2022.

We thank the reviewer for their constructive comments. It seems the score changes proposed by the reviewer might not have been reflected in the main review. We wanted to double check if this was the case from their end.

Thank you for taking the time to review our work. We are happy to hear that the reviewer found our work novel and easy to follow. We address the reviewer's questions and concerns below:

A bit more summary of baseline methods would better show which part of the paper is novel, while which was already discussed.

We have added more information about the baselines in Appendix B. To reiterate, our method identifies the affine nature of Bellman Equations using which it presents a method to learn representations for any MDP without imposing additional structure (e.g. HILP uses representations only suitable for goal reaching, and Laplacian representations are dependent on the policy used to collect the data) and avoids the instability of Bellman optimality backups (used in FB) that suffer from overestimation in Deep RL by only requiring Bellman evaluation backups.

In Line 368-370, some procedure to recover the optimal policy in the setting of continuous spaces is provided. However, the former theoretical results are based on proof in the setting of finite discrete spaces.

We clarify that our proposed method is not limited to discrete state or action spaces and extends to continuous domains as well. The matrix equations can be written using functional operators, and the linearity of the equation is still maintained. We have updated the paper in Section 4 to reflect this more clearly. We have also added experiments in continuous DM-control domains across 16 tasks to validate that PSM indeed scales to continuous domains.

As implied by Equation (10), the final step is to solve the dual problem, some theoretical analysis of the computational cost could help.

After our core pretraining phases of Sections 5.1 and 5.2, finding an optimal policy given any reward function amounts to solving a constrained linear program of Equation 10. The objective in Equation 10 is linear with affine constraints (thereby, a convex optimization problem). These classes of problems are well studied and guaranteed to converge to optimal solutions with a sublinear rate. More importantly, Equation 10 reduces the hard non-convex RL problem that needs to be solved for each reward function to a simple convex optimization for each reward function.

Punctuation and Matrix definition in Proof of Lemma 4.1

We thank the reviewer for pointing this out. We have corrected the punctuation and added the matrix definition in the paper.

In Section 4, toy example, there is a reference to the derivation of Equation (5), which does not appear in the Appendix.

Thanks for catching this. We have updated the paper to point to the derivation of Eq.5 in the supplementary material.

Without more explicit explanation of , currently Theorem 4.4. is not easily accessible. Some definitions would help.

We have added more definitions in Theorem 4.4 to make it easy to understand.

Equation (7) seems to be important in forming the method, however, model densities and other notations used there are not defined.

We have updated the paper in general to increase clarity. The definition for model densities is given in Section 5.1, first paragraph before using them in Equation (7). Please let us know if there are any other notations that are unclear.

Although the abstract implicitly claims to not make assumptions on the MDP structure, according to both theories and experiments, discrete state and action spaces are focused on in this paper

We apologize for the confusion. Our proposed method is not limited to discrete state or action spaces and extends to continuous domains as well. We have updated the paper to make this more clear. Furthermore, we have added evaluation on complex continuous state-action environments, demonstrating that PSM leads to significant improvement in these domains.

In the procedure in Section 5.3, is the reward function assumed to be known?

Yes, sections 5.1 and 5.2 form the core algorithm where we pre-train representations without assuming a reward function. Then, when a user provides us with a task (in the form of reward function), our fast inference procedure can quickly give an optimal policy by solving a constrained linear program.

Are Lemma 4.1 and Theorem 4.2 proposed by the paper or in former literature?

We have made minor modifications in the paper (highlighted in red) to make this clear. Both lemma 4.1 and theorem 4.2 are contributions of this work and novel to the best of our knowledge. To increase clarity, we have restructured these into Theorem 4.1 and Corollary 4.2.

The changes to the paper are marked in red. Along with increased evaluation, we have added detailed experimental details and made the writing more clear. Please let us know if you have any further questions. We hope the reviewer can reassess our work in light of these clarifications.

We thank the reviewer for the thoughtful review. We are glad to see that the reviewer found our method technically sound. We address the reviewer’s concerns below:

References for world models in reinforcement learning

We thank the reviewer for pointing this out and the reference has been fixed now.

While this simplifies optimization, it may also impose limitations. A discrete policy codebook can only represent a finite subset of the policy space

We clarify that while we use a discrete codebook to generate the policies, the number of codes can be unlimited. In fact, given an unbiased random code generator, we can show that our policies encapsulate all possible deterministic policies as the infinite number of codes and for practical settings, the discrete codebook provides an unbiased estimator of the expectation w.r.t sampled policies.

The experimental setup is fairly limited, testing only two basic environments (Grid World and FetchReach), with relatively simple tasks.

We have now included an evaluation of 16 continuous control tasks from the DM-control benchmark. These environments are typical for evaluating RL algorithms for continuous state-action spaces and PSM scales without difficulties in these domains.

The changes to the paper are marked in red. Along with increased evaluation, we have added detailed experimental details and made the writing more clear. Please let us know if you have any further questions. We hope that the reviewer can reevaluate our work in light of these clarifications.

We have summarized the improvements to the paper based on the reviewers' comments as the General Response. We have further added the following additional experiments in Appendix C:

a. Ablation on the dimensionality of affine space: We observe PSM performance to be stable across dimension sizes 32 to 256.
b. Quantitative results on grid-world and discrete maze environments: PSM outperforms the baselines.

As the discussion period is ending soon, it would be great to hear from the reviewer as we have addressed their concerns about the draft. Please let us know if you have any further questions. We would appreciate it if the reviewer could reassess our work in light of these clarifications and additional empirical results.

We thank the reviewer for the detailed review. We are happy to see that the reviewer found the idea of leveraging the property that visitations form affine sets novel, theoretically sound, and to have a potentially significant practical impact.

We address the perceived weaknesses and questions raised by the reviewer here:

The tasks chosen in the experiments section are quite simplistic, and it is not clear how the method would hold in more complex environments (e.g. with large state space or a more complex transition function).

The PSM loss function directly extends to continuous state-action spaces and we have now added experiments on the DM-control tasks, which represent a set of continuous control environments popular for benchmarking continuous RL agents. We have added evaluation for a variety of tasks defined by general reward functions; as opposed to goal-based tasks.

The technical details of the proposed algorithm are missing, e.g. how is parametrized, what optimizer is used, etc. Further, the technical details of the benchmarks are also missing, e.g. what is the horizon and reward function of each task (without these it is hard to assess the quality of the success rate), etc.

We thank the reviewer for pointing this out. We have now added a detailed summary of all the hyperparameters and experimental details in the appendix.

The limitations of the proposed approach were not discussed in the paper.

We have added a discussion about the limitations of our work in the paper in the conclusion section. The major limitation is that it is hard to ascertain a good size for dimensionality of basis functions. Complex MDPs with highly non-linear and non-smooth transition dynamics might require a larger dimensionality but a larger dimensionality than needed might lead to larger compute requirements and training time. An interesting future direction would be to investigate if ‘validation’ can help to determine the minimal dimensionality: running offline RL with the learned representations across a set of reward functions and choosing the dimension size that gives the highest average Q-values.

A high-level flow for Sections 4 and 5 could have helped the readability of the paper

Thanks for this suggestion. We have added a high-level flow for Sections 4 and 5.

A theoretical discussion and analysis of the dimensionality of the affine space is missing in Section 4 (i.e. what affects ).

We are currently running ablation experiments on how the dimensionality of affine spaces affects the performance of PSM. We expect to update this before the discussion period ends. A complete theoretical analysis of how the dimensionality of affine spaces affects performance is an interesting future direction and out of the scope of the current paper.

TYPOS

We have addressed this now, thank you!

Are there any practical benefits of one approach or another (two player vs single player)? I would suggest incorporating a discussion on this into the paper or adding some sort of an ablation study on that to the empirical experiments.

Two-player optimization is known to suffer from training instabilities when compared to single-player optimization (e.g. GAN vs Diffusion models, IRL vs. Behavior cloning, etc). Our initial attempts to train using a two-player setup led to performance collapse, and thus, we chose a more stable, principled and performant alternative for presenting our method.

How many seeds were used in the experiments in the empirical study?

We used 5 seeds for our continuous experiments and 3 seeds for our discrete experiments in Section 7 and have updated the paper to make this more clear.

You discussed the quality of the Q function based on the visualizations in lines 440-444. Did you also perform a quantitative study on this (e.g. by comparing to the oracle optimal Q function)?

We are running the quantitative evaluations on gridworlds, and the results will be updated before the discussion period ends.

Did you use the same offline dataset mentioned in the paper for all baselines?

Yes, we have now added a thorough explanation of implementation details in the appendix to make this clear.

Did you perform a hyperparameter tuning procedure for your algorithm and the baselines? If so, can you elaborate on this?

We have added the hyperparameters used for each method in the appendix. We kept the base hyperparameters the same (representation dimensionality, neural network capacity, inference samples, etc.) to ensure fair comparison but used the tuned hyperparameters for algorithm-specific hyperparameters (e.g. PSM inference learning rate). All hyperparameters can be found in Appendix B.3.1.

In the robotic manipulation task (Section 6.2), how did you handle a continuous state space? Doesn't that make the basis function continuous?

We clarify that PSM learns the basis of successor measures or successor visitation distributions, which are made up of continuous quantities. Thus, the PSM algorithm is invariant to discrete or continuous state or action spaces.

Please let us know if you have any more questions. We hope the reviewer can reassess our work in light of these clarifications.

We have summarized the improvements to the paper based on the reviewers' comments as the General Response. We have further added the following additional experiments in Appendix C as requested by the reviewer:

a. Ablation on the dimensionality of affine space: We observe PSM performance to be stable across dimension sizes 32 to 256.
b. Quantitative results on grid-world and discrete maze environments: PSM outperforms the baselines.

As the discussion period is ending soon, it would be great to hear from the reviewer as we have addressed their concerns about the draft. Please let us know if there are any further questions. We would appreciate it if the reviewer could reassess our work in light of these clarifications and additional empirical results.

Thank you for increasing the score and helping us with improving the paper!

As a matter of public record, we are responding to the AC meta-review as there were several new points brought up. We summarize the new reasons the AC provided for their reject decision in spite of reviewer ratings (8, 8, 6, 5) and detail our response to each below.

The AC mentioned that a large number of experiments were added after the submission deadline. These experiments were added in response to the questions asked by the reviewers. All reviewers wanted to see the performance of our method on a continuous domain so we provided a comparison of our method with the baselines on this domain. We believed that our research contribution in terms of the idea of representing all possible solutions of RL was already strong enough containing theoretical and empirical evidence to be a worthy submission, which reviewers clearly agreed with due to their change in scores.

The AC mentions that there are several major errors in experiments that prevent accurate interpretation of the results. The AC believes that our method lacks statistical testing with other methods. Additionally, the AC points out that we present our performance comparisons with a specific set of hyperparams. We believe we followed the standard practice to compare the mean performance across several seeds which was also used by the baselines [1, 2, 3]. While statistical testing across several hyperparams is a stronger evaluation, we submit that we shouldn’t be penalized for following the standard practice which is also followed by several of the papers accepted at ICLR 2025.

The AC points out some clarity concerns in Theorem 4.4 of our paper. We thank the AC for their suggestions to further improve the clarity of the Theorem. We had modified the Theorem statement following suggestions by a reviewer but neglected to modify the corresponding proof in the appendix, which is easily corrected. We can further modify the statement and the proof in the camera ready version. We would like to point out that while Theorem 4.4 is important to our discussion, it is not the main contribution of the paper. The AC claims that our statement in the introduction about eigenvectors from a random policy not being able to represent any value function is false. The AC mentions that if the graph laplacian is full rank, it is possible to represent any policy’s value function. Our point is that for all practical purposes, it is generally not possible to get a full rank graph laplacian or to have a random policy that covers the entire state space. Hence in practical settings, it is impossible for eigenvectors from a random policy to be enough to cover all possible value functions.

The AC makes an incorrect statement in their final remark by saying “there were limited back and forth with reviewers” (and then nonetheless nothing that “several reviewers adjusted their scores positively”). In fact, we submitted a total of 28 comments in the review process. As you can see from our responses, we were very prompt in replying to reviewers.

We thank the AC for thoroughly evaluating our paper and pointing out typos and concerns in clarity. But, the points raised by the AC are relatively minor, irrelevant to the main contribution, and/or easily addressable in camera-ready.

[1]: Touati, Ahmed, Jérémy Rapin, and Yann Ollivier. "Does Zero-Shot Reinforcement Learning Exist?." The Eleventh International Conference on Learning Representations. [2]: Park, Seohong, Tobias Kreiman, and Sergey Levine. "Foundation Policies with Hilbert Representations." Forty-first International Conference on Machine Learning. [3]: Wu, Yifan, George Tucker, and Ofir Nachum. "The Laplacian in RL: Learning Representations with Efficient Approximations." International Conference on Learning Representations.