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

Authors: Siddhant Agarwal, Harshit Sikchi, Peter Stone, Amy Zhang
Primary area: reinforcement learning
Keywords: Zero-Shot Reinforcement Learning, Representation Learning, Unsupervised RL

Abstract:
Having explored an environment, intelligent agents should be able to transfer their knowledge to most downstream tasks within that environment. 
Referred to as ``zero-shot learning," this ability remains elusive for general-purpose reinforcement learning algorithms.  While recent works have attempted to produce zero-shot RL agents, they make assumptions about the nature of the tasks or the structure of the MDP. We present \emph{Proto Successor Measure}: the basis set for all possible solutions of Reinforcement Learning in a dynamical system. We provably show that any possible policy can be represented using an affine combination of these policy independent basis functions. Given a reward function at test time, we simply need to find the right set of linear weights to combine these basis corresponding to the optimal policy.  We derive a practical algorithm to learn these basis functions using only interaction data from the environment and show that our approach can produce the optimal policy at test time for any given reward function without additional environmental interactions.

Decision: reject
Reviewer scores: 8, 6, 8, 5

Reviews and discussion:
<thread>
creation date:
1729683804564

thread_type:
official_review

<thread_object>
creator: reviewer_nqqb

<content>
summary:
The authors introduce a Proto Successor Measure (PSM), which is a method to solve RL problems with a focus on achieving zero-shot learning. PSM works by creating a basis set for all possible solutions in an environment by decomposing the state-action visitation distributions, using linear combinations of policy-independent basis functions. This allows an agent to compute the optimal policy for any given task by adjusting the weights associated with the basis functions.

strengths:
1: The authors claim: "*We show that any solution of the Bellman Flow constraint for the visitation distribution
can be represented as a linear combination of policy-independent basis functions and a bias*". In Chapter 4 and the appendix, the mathematical reasoning to support this claim looks sufficient and is well-written. However, I do not know whether the authors are the first to prove this. If so, it would be beneficial to point this out in the paper, preferably at the end of the introduction.

Other strengths:
- Interesting idea, and a nice visual comparison to baselines in a maze environment.
- The paper is easy to read through.

weaknesses:
1: 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.

2: The claim: "*We show that PSM can produce the optimal Q function and the optimal
policy for any goal conditioned task in a number of environments outperforming prior baselines.*" is not sufficiently proven. First of all, only a maze environment and Fetchreach are used, which does not correlate to "a number of environments". Second, the authors claim to outperform prior baselines. However, performance is only compared in the Fetchreach environment. 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.

3: Evaluation: No other evaluations or ablations can be found in the very limited Appendix section.

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

Other weaknesses:
- 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.
- References: The references are organized very poorly:
    - Line 535, Line 540, Line 543 , Line 576, Line 579, Line 581, Line 591, Line 593, Line 599: All these references have no venue.
    - Line 611 & Line 616: These are two double references of exactly the same papers.
    - For a lot of references that have been published in conferences, only the Arxiv reference is given.
    - This correlates with my statement in weakness 1, where it just feels that this paper has not been carefully prepared but has been submitted prematurely.

[Ahmed Touati & Yann Olivier, 2021] - Learning One Representation to Optimize All Rewards, NeurIPS 2021.

questions:
- Why are there no ablations, and no other quantitative evaluations other than the Fetchreach environment?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730422112341

thread_type:
official_review

<thread_object>
creator: reviewer_nimh

<content>
summary:
The paper is considering the 'zero-shot learning' problem, where the RL agent learns to understand a given environment and is expected to solve downstream tasks with different reward functions. The authors first show that any valid policy, by using its visitation distribution and successor measure, can be represented by some affine combination of basis vectors (in the discrete settings). And a corresponding algorithm is proposed to learn such basis vectors and then solve the optimization problem through the form of a  Lagrangian dual. A couple of experiments are provided to show its advantage.

strengths:
> **Originality**
- The paper proposed a novel perspective to represent the set of all valid policies using their visitation-based probability measures, which further gives rise to a concise representation of the optimization problem in zero-shot RL.

> **Quality**
- Despite notational confusions, theoretical results provide some new insights for understanding an arbitrary policy.

> **Clarity**
- The presentation is mostly easy to follow, though some notations/theorem need more clarifications.

weaknesses:
> **Originality**
- A bit more summary of baseline methods would better show which part of the paper is novel, while which was already discussed.

> **Quality**
- 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.
- As implied by Equation (10), the final step is to solve the dual problem, some theoretical analysis of the computational cost could help.

> **Clarity**
- In Equation (1) - (3), a comma or $\times$ is missing between $\mathcal{S}$ and $\mathcal{A}$.
- In Section A.1, Proof of Lemma 4.1: first, the left bracket is in the wrong place in Line 659; and an explicit definition of matrix $S$ would help
- In Section 4, toy example, there is a reference to the derivation of Equation (5), which does not appear in the Appendix.
- Without more explicit explanation of $\Phi\^{vf}$, currently Theorem 4.4. is not easily accessible. Some definitions would help.
- Equation (7) seems to be important in forming the method, however, model densities and other notations used there are not defined.

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

questions:
- In the procedure in Section 5.3, is the reward function assumed to be known?
- Are Lemma 4.1 and Theorem 4.2 proposed by the paper or in former literature?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730696856482

thread_type:
official_review

<thread_object>
creator: reviewer_polf

<content>
summary:
This paper presents Proto Successor Measure (PSM), a zero-shot RL method that compresses any MDP to allow for optimal policy inference for any reward function without additional environmental interactions. The method seems technically sound.  But experiments are limited.

strengths:
1. The proposed Proto Successor Measure (PSM) method combines policy-independent basis functions with biases, enabling it to generate new policies from offline data without additional environment interactions, and allowing it to adapt to different tasks.
2. The paper introduces a policy-independent basis function set, Φ, which can be shared across different tasks and policies. Since these basis functions are independent of specific policies, new policies can be generated simply by adjusting the combination weights. This reduces the need for retraining on new tasks, enhancing generalization capability and computational efficiency.

weaknesses:
1. The paper cites "Planning with Diffusion for Flexible Behavior Synthesis" for its world model, but this work primarily focuses on using diffusion models for trajectory planning and is not directly related to world models in reinforcement learning. Relevant world-model references could include "Dream to Control: Learning Behaviors by Latent Imagination."
2. The paper represents policies as a compact discrete codebook, approximating an infinitely large policy space through a discrete set. While this simplifies optimization, it may also impose limitations. A discrete policy codebook can only represent a finite subset of the policy space, potentially leaving some strategies uncovered and leading to suboptimal results in certain tasks.
3. The experimental setup is fairly limited, testing only two basic environments (Grid World and FetchReach), with relatively simple tasks. There is a lack of more complex, high-dimensional environments (such as robot control tasks in Meta-World or Mujoco) to demonstrate PSM's generalization capabilities and applicability. This limitation may affect the persuasiveness of the method's performance in a broader range of environments.

questions:
See previous section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730720239594

thread_type:
official_review

<thread_object>
creator: reviewer_mxvw

<content>
summary:
In the paper, the authors present a novel approach to unsupervised reinforcement learning. They first present a theoretical foundation for the algorithm by showing that any state-action visitation distribution and any successor measure $M^{\pi}$ in a single MDP form an affine set. This motivates their approach - first learn the basis of the affine set (in an unsupervised manner, without a specific policy or reward), and then, given the affine space, the only optimization left given a downstream task is to find the optimal linear representation in the space. The authors present a practical algorithm to achieve both steps and evaluate it on two simple benchmarks and a robotic manipulation benchmark while comparing to two common unsupervised-rl baselines.

strengths:
1. As far as I know, the authors are the first to make the observation that all state-action visitation distributions in an MDP can be represented as an affine set. This can have a potential impact on the field of unsupervised-rl in particular and other fields in rl in general. 

2. The theoretical claims in papers are well supported and presented in a precise and clear way

3.  The suggested algorithm can potentially leverage large offline datasets to effectively learn representations for downstream zero-shot tasks, which can have a significant practical impact.

weaknesses:
1. 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). 

2. The technical details of the proposed algorithm are missing, e.g. how $\Phi$ 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. 

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

4. While the paper is well-written and the ideas are presented clearly, it is hard to keep track of its high-level flow. A high-level flow for Sections 4 and 5 could have helped the readability of the paper

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

questions:
1. Typos:

    a. Line 432: toform -> to form
    
    b. Line 279 and Line 257: Add figure numbers

2. In line 267 you mention the derivation for the basis vector is in the supplementary material. I did not find such a derivation, is it missing? If so please add one and add an exact reference to it.  

3. At the end of Section 5.1 you presented a two-player game optimization strategy and suggested simplifying it with the single-player game presented in Section 5.2. Are there any practical benefits of one approach or another? 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. 

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

5. 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)?   

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

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

8. In the robotic manipulation task (Section 6.2), how did you handle a continuous state space? Doesn't that make the basis function continuous?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731778237018

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank all reviewers for taking the time to review our work. We are encouraged to see reviewers point out: the importance of recognizing that visitations form an affine set and its potential impact to the field of unsupervised RL (Reviewer mxVW), that the paper has sound theory with theoretical justification (Reviewer mxVW, PoLf, NimH, nQqb) along with writing that was clean and easy to read (Reviewer mxVW, NimH, nQqb). A common issue raised by reviewers was limited evaluation to support the theory (PoLf, nQqb). We have addressed this limitation in the rebuttal.

The main changes to the paper are highlighted in red in the PDF. While we respond to each reviewer individually, we outline the general improvements to the paper in this comment:

1. **Continuous control benchmark**: We evaluate PSM across 16 tasks in 4 continuous control domains. Along with the previously presented results in discrete environments, our evaluation demonstrates significant benefits from the proposed method PSM. PSM leads to 11.5% improvement over the best-performing representation learning method (FB).


2. **Writing improvements**: We incorporated advice from the reviewers to correct typos, increase readability, and use the correct references that point to published version of arXiv papers. We have further added a section discussing theoretical connections between PSM and successor features


3. **Complete experimental details in appendix**: We have added complete details for reproducing our experiments in appendix section B.

4. **Ablation on latent dimension**: We have added ablations on dimensionality of affine space in Appendix C.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a method for learning features from a set of data that are useful for quickly identifying an optimal policy for a new task (new reward function). The central component of the method is to learn successor features because leveraging the Bellman flow constraints on the state-action visitation distribution forms a space where the value function is a linear successor measure. 

The reviewers were generally positive about the paper’s intent and proposed contributions. However, there are significant issues of clarity; a large number of experiments added after the paper deadline that could not be properly reviewed, and some serious flaws in the experiments that lead to ambiguous conclusions. 

I detail issues not discussed by the reviewers below. 

In reference to the paper’s clarification, there are many typos, undefined portions, and some false statements. 
In the intro, the line “However, the eigenvectors from a random policy cannot represent all value functions.” is false. If the graph laplacian is full rank, it can represent any policy’s value function. Thus this produces a misleading characterization of the motivation for the paper. 

In Theorem 4.4 (main paper) the function r(s+) is not defined. w^pi is used in both definitions of the value function, but are dependent on the basis functions \Phi and \Phi^{vf}. These basis functions also lack clear definitions. There is a “.” in the expression for V^\pi, which should not be there. 

In Theorem 4.4 (appendix) the theorem statement does not match the main body. The proof also makes a special assumption that is never stated nor proved to ever be true. The proof is not easy to follow and has a term \Beta^\pi_i, which is never used. It also makes unclear references “The two equations match,” it is unclear which equations these are referencing. 

These errors and lack of clarity (in addition to the ones noted by the reviewers) make it difficult to accurately review the paper and be sure there are not other mistakes hidden in the paper. 

In terms of the experiments, there are several major errors that prevent accurate interpretation of the results. First, all quantitative results are based on three or five random seeds, and no statistical testing is given to make comparisons. Despite this, the results are bolded to indicate that the proposed method is superior to the others. Second, there is no control over how the hyperparameters were tuned or selected. As a result, the comparison could only be valid for one specific set of hyperparameters, which is of little interest and likely to be biased to favor the new method over old ones. At most, what can be learned from the experiments is that the new method works and perhaps works well with the chosen hyperparameters, but it is unclear if it represents an improvement over existing methods. 

As a result of the above, I cannot recommend the paper for acceptance. However, I do recommend the authors address these issues, improve the clarity, and resubmit.