Proto Successor Measure: Representing the Behavior Space of an RL Agent

We thank the reviewer for providing detailed feedback on our paper. We would like to address the concerns raised by the reviewer:

At best, further samples are needed from the environment to infer a reward function. This could be problematic especially in sparse environments. :

We agree with the reviewer’s observation that practically sparse reward environments can be problematic. Mathematically, our method can be used to obtain the near-optimal policy for any reward function. But in practice, the quantity $\sum_{s’} \phi(s, a, s’)r(s’)$ is approximated by $\mathbb{E}_{s’ \sim \rho}[\phi(s, a, s’) r(s’)]$. This restricts the knowledge of reward function to only samples from the function, which could affect the optimality in some situations. We will soften the claim in the camera-ready version.

Theorem 4.4: The notation is somewhat confusing for me here. Fixing typography can also help readability.:

We thank the reviewer for pointing out that the Theorem statement might be confusing. We shall rewrite the theorem statement in the camera ready version. Through this theorem, in line with reviewer’s understanding we wished to portray that spanning successor measures using a basis covers a larger set of RL solutions than spanning value functions using a basis.

The proof seems to follow with a straightforward rearrangement of terms, but it is a bit opaque (to me):

The proof establishes that for every instance of a $V^\pi \in span{\phi^{vf}}$, there exists an instance of $V^\pi$ in $\span{\Phi}r$ for some reward. While the reverse is not always true. $\beta^\pi$ is used to denote the weights of the linear combination of \phi^{vf} while representing $V^\pi$ and $k_i = \sum_{s’}r(s’)$ is used for algebraic manipulation. We appreciate the reviewer’s concern and will make the proof clearer in the camera-ready version by adding explanations for the different algebraic manipulations used.

Fig 4: What are update steps? seeds is quite small -- adding more could strengthen the statistics.:

Update steps means the number of learning updates made to the bases and bias. $\phi$ and $b$ were evaluated on their zero-shot performance after every fixed number of training steps.

We appreciate the reviewer for pointing this out and will add more seeds during the camera ready version.

Table 4: Do you know how such errors scale with size of maze?:

There are 50 and 68 cells respectively in the grid world and four room maze. So the % of wrong predictions for PSM are 4.1% and 16.97%. This is in comparison to 29.06% and 42.35% for FB and 38.56% and 56.57% for Laplace.

How is there any "smoothness" between goals in the DMC tasks? They seem quite distinct, so how is there a generalization effect among tasks?

We are not sure if we understand the reviewer’s question - Most of the tasks we evaluate for DMC use dense rewards as opposed to goals. These tasks are the standard ones that prior methods (Touati et. al., 2023; Park et. al. 2024) used for evaluating unsupervised RL performance. PSM uses the unlabelled transition dataset to learn all possible behaviors (stitched trajectories). Near optimal policies for each of these tasks correspond to a particular stitching.

Could you further clarify the use/meaning of the bias vector and how its presence departs from typical SF?:

The bias vector comes from the fact that the Bellman equation is affine and not linear. The significance being that a zero visitation/successor measure is not permitted which might be permitted in SF. Note that it is possible to recover the bias by representing successor measures as $\phi w$ (if the last coordinate of $w$ being learnt to be 1). But, the bias term forces a mathematically correct representation. What truly departs from SF is the absence of the assumption of the linear mapping between rewards and policy embeddings.

Is there any generality lost in the linearity of the framework? What are possible downsides/alternatives/missing generalizations/adversarial reward functions to keep in mind?:

The PSM objective ideally will represent all possible successor measures contrary to SF-methods which focus on learning successor measures optimal for a particular class of reward functions (linear-in-state-representation). This means they will fail to capture the successor measure if it was not optimal for a certain reward function belonging to the class. Second, assuming linearity between rewards and zero-shot policy embeddings, fundamentally disallowing many-to-many mapping between policy and rewards.

Can you provide an experiment showing the result of small datasets/dataset scaling? :

We run PSM on a reduced dataset (by selecting only 10%) of the trajectories in the RND dataset for walker and cheetah. While the performance on walker remains 656.08 (full dataset performance: 689.07), the performance on cheetah dips to 211.53 (full dataset performance: 626.01). We will add detailed experiments in the camera ready version.

We thank the reviewer for taking time to review our paper and providing detailed feedback. We would like to address the concerns of the reviewer below:

I think the "mathematically complete" mention is a bit exaggerated here. The paper would require a much more quantitative analysis to deserve that mention.:

By “mathematically complete,” we mean that the representation learning framework is derived without making any assumptions on (1) dynamics, (2) reward functions, or (3) the complexity of the state or action space. Since the practical algorithm has some approximations to make the framework more tractable, starting with the low-rank approximation of the bases space itself, we thank the reviewer for pointing this out and understand it might be construed to mean something stronger.Wee shall soften and clarify this language in the camera-ready version.

I am not sure think PSM really compares with FB, Laplacian and the likes:

We respectfully disagree with the reviewer that PSM does not really compare with FB, Laplacian. In all our experiments, the respective functions, ($\phi, b$ for PSM or $F, B$ for FB) are represented using neural networks with a similar number of parameters, thus working with similar hypothesis spaces. While PSM requires training the bases with networks taking in (s, a, s’), FB requires training two networks, one with input (s, a, z) and the other with (s). The dimension of z is similar to (often more than) s (128 in our experiments). This observation, in some ways, makes FB require more parameters than PSM. But, in any case, FB and other SF based approaches like Laplace and HILP seem to be the most relevant prior work in terms of zero-shot RL that PSM can compare against.

Typos:

We thank the reviewer for pointing out the typos. We shall correct them in the camera-ready version.

I think the theoretical content is quite simple so I would suggest putting the focus on the experimental work.:

We are glad to see the reviewer likes our observation that successor measures span an affine space. While the idea may seem simple in hindsight, it has (to the best of our knowledge) remained unexplored in the RL literature. Hence we decided to spend more time building towards insights we can obtain from the affine-space idea of successor measures - for instance we are able to show that with the same basis dimensions, we can span a larger space of values than methods which learn a basis space across value functions.

The paper puts the emphasis on state-action distributions and regards the successor measure as a secondary. I believe this should be the other way around, the successor measure is the central object and state-action distributions a byproduct of it.:

We are aligned with the reviewer's observation that successor measures are the main object we end up learning to represent. For unsupervised learning, it is preferable to learn to represent successor measures as they capture more information that visitation distributions. Our choice of starting with visitation distributions was motivated by the large number of works that explore the idea of RL as a linear program and goes to show the same idea can be used for unsupervised representation learning. The familiarity of readers with that literature motivated us to start from visitation distribution to facilitate understanding. We will clarify the writing to place more emphasis on the successor measure.

Is it true as I think that PSM and FB do not play in the same league exactly, since the features of PSM are functions of (s, a, s’)? Is it really possible to learn PSM as efficiently as FB?:

As discussed earlier, FB and PSM are parameterized with neural networks with a similar number of parameters, in fact, depending on the sizes of s and z, FB may be parameterized with a larger network. Empirically it has been observed that FB is difficult to stabilize as it is optimizing for a moving reward function ($B^{-1} z$) while PSM does not face this issue as it does not tie the policy latents to reward functions. Moreover, FB uses the policy optimized for the corresponding reward (through Q function maximization) to sample actions during its training. This means that FB is susceptible to overestimation issues occurring due to Q function maximization.

We appreciate the reviewer for providing detailed feedback on our paper. We would like to clarify all the questions raised by the reviewer:

The inference for the linear weights involves solving a linear program with constraints. This is harder to solve than linear regression in prior works. :

Yes, while we agree that the full PSM objective requires solving a constraint LP, which can be harder than simple linear regression as in some of the prior work, other methods relying on linear regression rely on the assumption that policy embeddings z have a linear relation to reward functions. Enforcing this assumption in PSM, as we have shown in Theorem 6.3 (Line 318-328 right column), PSM can also enable inference through linear regression.

PSM does not directly produce policies. It produces the successor measure corresponding to the optimal policy for the downstream reward. To back out a policy for a continuous domain, one needs to recover the Q value and then perform several iterations of policy optimization to take actions that maximize the Q values. In some prior work, the definition of zero-shot RL is no additional policy optimization. PSM would violate this strict definition of zero-shot RL. :

Following our explanation in comment 1 above, the full objective of PSM does require additional optimization to obtain policy from successor measures, but these optimizations are much easier than solving an RL problem and similar to a policy evaluation problem. With the additional assumption of rewards linear in state features obtained by PSM, we can enable true zero-shot inference without additional optimization. The key benefit of PSM lies in its flexibility to accurately represent all behaviors of an agent and with an additional assumption can also outperform prior works in zero-shot RL.

The paper does not address the data assumption. In principle, the dataset needs to have full coverage of the state and action spaces to learn a basis set that spans the entire space of successor measures:

The capabilities of any unsupervised learning algorithm is intrinsically tied to the dataset it is learning from. PSM and its unsupervised RL counterparts FB and HILP all are limited by dataset but PSM aims to learn to represent a more diverse range of skills as a virtue of its objective. But this very property can be beneficial in practice: When using real world dataset to learn skills we would avoid focusing on skills that are irrelevant in the real world. With complete coverage of the space, PSM will learn the basis that attempts to span the entire space of successor measures.

Typos:

We thank the reviewer for pointing out these typos. We shall fix them in the camera-ready version.

How do you solve (4) in continuous space? Do you sample (s, a) from a dataset?

In both continuous and discrete spaces, we sample (s, a) from the dataset.

How do you determine the dimensionality of the basis set? Can you provide an ablation of the dimensionality of the bases?

We have provided an ablation for the dimensionality of the bases in Appendix C.1. The dimensionality of the bases can potentially vary depending on the dynamics. For instance, If the dynamics is very regular (symmetric and Lipschitz as a lot of the real world domains are), the dimensionality may be way less than the mathematical limit of S x A. An example can be constructed similar to the one in Section B.2 in [1].

[1]: Touati, Ahmed, and Yann Ollivier. "Learning one representation to optimize all rewards." Advances in Neural Information Processing Systems 34 (2021): 13-23.

In hyperparameters (Table 2), why does w have 3 layers? Is it represented by a neural network?

During pretraining i.e., obtaining the bases of PSM, we represent w as a function of the policy representation: the discrete code c. As a result, we represent w as a neural network which is a mapping from c to $\mathbb{R}^d$.

We thank the reviewer for their detailed feedback. We would like to address all the concerns of the reviewer below:

It would be good to compare a more diverse set of reward functions for each environment.:

We appreciate the reviewer’s suggestion and agree that our original evaluation using four tasks per environment, though consistent with prior works, could be further expanded for stronger claims. To directly address this concern, we have conducted additional experiments on a significantly larger and more diverse set of reward functions (10+ tasks) in the Walker environment. Our extended results confirm that PSM consistently outperforms or matches the baselines. The mean performance on these 10+ tasks are (across 5 seeds), PSM: 523 +- 48.86, FB: 516.95 +- 45.61 and Laplace: 250.43 +- 30.58. We shall include detailed results with diverse tasks for each of the domains in the camera ready.

Comparison with successor feature methods:

We would like to point out that the baselines that we considered are in fact SF-based approaches. “Laplace” uses Laplace eigenvectors as state features while HILP has its own objective for learning state features. Both of them learn Successor Features for these corresponding state features. Moreover, FB is also built off an SF-based approach. We acknowledge the reviewer and will make it clearer in the paper. Additionally, we are adding experiments on a couple of other SF-based approaches (best performant as per (Touati et. al., 2023): one that uses one-step dynamics predictability (method named “Trans”) and one that uses SVD decomposition of successor representation (method named “SVD”). The mean performance for these baselines across all DMC tasks are (across 5 seeds): Trans: 522.10 and SVD: 524.96. PSM outperforms them by a significant margin with a mean performance: 607.60. We will add the detailed results in the camera ready.

However, this is quite different from the results provided in the original FB paper. It is unclear where the differences come from. :

We thank the reviewer for noting this difference. We would like to clarify that the original FB paper does not report quantitative performance metrics comprehensively on the four-room environment but rather highlights qualitative successes with select examples.

Additionally, to ensure fair comparison with FB, we deliberately modified its training strategy. The original FB method assumes that tasks are goal-conditioned, implicitly biasing the training to a smaller set of reward functions. Since PSM makes no such assumption, we adapted FB training to uniformly sample reward functions to ensure an apples-to-apples comparison. This explains the observed discrepancy in FB performance. We will clearly document these modifications in the final version. We shall include the results for the biased sampling in FB in the final version.

While Section 6 points out that SF can be viewed as a special case of PSM, it is not well explained why it is beneficial to consider PSM in general. :

We appreciate the reviewer highlighting the need for clearer explanation of the advantages of PSM over traditional SF approaches. SF-based methods rely on two restrictive assumptions:

(1) Sufficient Feature Representation: SF assumes either predefined state features or features learned through an auxiliary objective, presuming these features adequately span all possible reward functions which limits representational expressiveness.

(2) Linear Mapping between Rewards and Policies: SF-based methods condition policies linearly on feature weights, assuming a one-to-one mapping between optimal policies and rewards. Practically, the mapping is many-to-many; multiple rewards can share the same optimal policy and vice versa, making such linearity limiting.

In contrast, PSM avoids both assumptions and directly aiming to accurately capture all feasible visitation distributions.

It remains unclear to me why PSM is a better approach than the FB representation:

(1) FB inherits limitations from SF-based approaches described above due to a similar linear mapping. (2) During training, FB samples from policies that are optimized (through Q function maximization) for the corresponding reward ($B^{-1} z$). This maximization step is prone to overestimation (3) FB is unstable due to optimizes for a changing reward function as B gets updated.

The statement of Theorem 4.4 is quite hard to parse, and the presentation of it can be improved. This is partially because the notations are quite confusing. :

We thank the reviewer for pointing this out. We agree and will revise the statement and the notation more clearly in Theorem 4.4 in the final version to improve clarity.

An ablation study on the need for the discrete policy codebook would be helpful.:

The discrete codebook is an essential component of our method that makes the PSM training tractable. We will add an ablation on the size of the codebook in the final version.