Proto Successor Measure: Representing the Behavior Space of an RL Agent
Authors: Siddhant Agarwal
Harshit Sikchi
Peter Stone
Amy Zhang
Venue: ICML 2025 poster (2025)
Keywords: Unsupervised Reinforcement Learning
Zero Shot Reinforcement Learning
Representation Learning
Primary area: reinforcement_learning
TL;DR: A zero shot RL method that represents all possible policies in the MDP

Abstract: Having explored an environment, intelligent agents should be able to transfer their knowledge to most downstream tasks within that environment without additional interactions. 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 *Proto Successor Measure*: the basis set for all possible behaviors of a Reinforcement Learning Agent in a dynamical system. We prove that any possible behavior (represented using visitation distributions) 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 bases corresponding to the optimal policy. We derive a practical algorithm to learn these basis functions using reward-free interaction data from the environment and show that our approach can produce the near-optimal policy at test time for any given reward function without additional environmental interactions. Project page: agarwalsiddhant10.github.io/projects/psm.html.

Proto Successor Measure: Representing the Behavior Space of an RL Agent
Siddhant Agarwal" Harshit Sikehi
Peter Stone 12 Amy Zhang
Abstract
to explore various routes figure out the most efficient way to get to the nearest supermarket place of work. But eventually, you could probably travel to new places effi- ciently the very first time you needed to get there. Like humans, intelligent agents should be able to infer the neces- sary information about the environment during exploration and use this experience for solving any downstream task efficiently. Reinforcement Learning (RL) algorithms have seen great success at finding sequence of decisions that optimally solves given task in the environment (Wurman et al., 2022; Fawzi et al., 2022). In RL settings, tasks are defined using reward functions with different tasks having their own optimal agent policy or behavior corresponding to the task reward. RL agents are usually trained for a given task (reward function) or on a distribution of related tasks; most RL agents do not generalize to solving any task, even in the same environment. While related machine learning fields like computer vision and natural language processing have shown success at zero-shot (Ramesh et al., 2021) and few-shot (Radford et al., 2021) adaptation to a wide range of downstream tasks, RL lags behind in such functional- ities. Unsupervised RL aims extract reusable informa- tion such skills (Eysenbach et al., 2019; Zahavy et al., 2023), representations (Ghosh et al., 2023; Ma et al., 2023), world-models (Janner et al., 2019; Hafner et al., 2020), or goal-reaching policies (Agarwal et al., 2024; Sikchi et al., 2024a) from the environment using aask-independent data to efficiently train RL agents for any task. Recent advances in unsupervised RL (Wu al., 2019; Touati & Ollivier, 20; Blier al., 2021b; Touati al., 2023) have shown some promise towards achieving zero-shot RL.
Having explored environment, intelligent agents should be able to transfer their knowledge to most downstream tasks within that environ- ment without additional interactions. Referred to as "zero-shot sive for learning", this ability remains elu- reinforcement learning algorithms. While recent works have attempted to produce zero-shot RL agents, they make as- sumptions about the nature of the tasks or the structure of the MDP. We present Proto Successor Measure: the basis set for all possible behaviors of a Reinforcement Learning Agent in a dynami- cal system. We prove that any possible behavior (represented using visitation distributions) can be represented using an affine combination of these policy-independent basis functions. Given re- ward function at test time, we simply need to find the right set of linear weights to combine these bases corresponding to the optimal policy. We derive practical algorithm learn these basis functions using reward-free interaction data from the environment and show that our approach can produce the near-optimal policy at test time for any given reward function without additional en- vironmental interactions. Project page: agarwal- siddhant 10t......ghhiii
1. Introduction
A wide variety of tasks can be defined within an environ- ment (or any dynamical system). For instance, in navigation environments, tasks can be defined to reach a goal, follow a path, reach a goal while avoiding certain states etc. Once familiar with an environment, humans have the wonderful ability to perform new tasks in that environment without any additional practice. For example, consider the last time you moved to new city. At first, you may have needed
Recently proposed pretraining algorithms (Stooke et al., 2021; Schwarzer et al., 2021b; Sermanet et al., 2017; 7; Nair et al.; Ma et al., 2023) use self-supervised learning to learn representations from large-scale data to facilitate few-shot RL but these representations are dependent on the policies used for collecting the data. These algorithms assume that the large scale data is collected from "good" policy demon- strating expert task solving behaviors.
Equal contribution 'The University of Texas Austin 2 "Sony AI. Correspondence to: Siddhant Agarwal <sid- d www.eesttxaascuuut. Harshit Sikchi <hsikehi utexaa.eduu> edu>>
Several prior works aim to achieve generalization in multi-task RL by building upon successor features (Dayan, 1993) which represent rewards as linear combination of state features. These methods have limited generalization capacity to unseen arbitrary tasks. Other works (Mahadevan,
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
Φw +
All Solutions
Environment
1
Valid Solutions
12
Task Reward:
Solutions Polytope
Figure 1. Method Overview: Visitation distributions corresponding any policy must obey the Bellman Flow constraint for the dynamical system. This means they must lie plane defined by Bellman Flow equation. Being a plane, can be represented using basis set & and a bias. All valid (non negative) visitation distributions lie within convex hull this plane. The boundary of this hull is defined using the non negativity constraints: 4ur + b 2 0. Each point within this convex hull corresponds to a visitation distribution for valid policy and is defined simply by the "coordinate" ur.
2005; Machado et al., 2017: 2018; Bellemare et al., 2019; Farebrother et al., 2023) represent value functions using eigenvectors of the graph Laplacian obtained from random policy to approximate the global basis of value functions. However, the eigenvectors from random policy usually cannot represent all value functions. In fact, we show that an alternative strategy of representing visitation distributions using set of basis functions covers a larger set of solutions than doing the same with value functions. Skill learning methods (Eysenbach et al., 2019; Park et al., 2024b; Eysenbach et al., 2022) view any policy combination of skills, but shown by (Eysenbach et al., 2022), these methods do not recover all possible skills from the MDP. Some recent work has attempted zero-shot RL by decomposing the representation of visitation distributions (Touati & Ollivier, 2021a; Touati et al., 2023), but they learn policy representations
in Figure 1, any visitation distribution that is a solution of the Bellman Flow for a given dynamical system lies plane defined using policy independent basis & and bias b. On this plane, only small convex region defines the valid (non- negative) visitations distributions. Any visitation distri- bution in this convex hull can be obtained simply using the "coordinates" w. We introduce Proto-Successor Measure, the set of basis functions and bias to represent any successor measure (a generalization over visitation distributions) in the MDP that can be learnt using reward-free interaction data. At test time, obtaining the optimal policy reduces to simply finding the linear weights to combine these basis vectors that maximize its dot-product with the user-specified reward. These basis vectors only depend on the state-action transition dynamics of the MDP, independent of the initial state distribution, reward, or policy, and can be thought to compactly represent the entire dynamics.
projection of the reward function which lead to loss of task-relevant information. We present stronger, more principled approach for representing any solution of RL in the MDP.
The contributions of our work are (1) novel, principled perspective on representation learning for Markov decision processes; (2) an efficient practical instantiation that reduces basis learning to single-player optimization; and (3) evaluations of a number of tasks demonstrating the capability of our learned representations to quickly infer near-optimal policies.
Any policy in the environment can be represented using visitation distributions or the distributions over states and actions that the agent visits when following a policy. We represent any possible visitation distri-
learn basis bution in the underlying environmental dynamics. We draw our inspiration from the linear programming view (Manne, 1960; Denardo, 1970; Nachum & Dai, 2020; Sikchi et al., 2024b) of reinforcement learning; the objective is to find the visitation distribution that maximizes the return (the dot-product of the visitation distribution and the reward) subject the Bellman Flow constraints. We show that any solution of the Bellman Flow constraint for the visitation distribution can be represented as a linear combination of iedeeeiindeenndent basis functions and bias. As shown
2. Related Work
Unsupervised Reinforcement Learning: Unsupervised RL generally refers to a broad class algorithms that use reward-free data to improve the efficiency of RL algorithms. We focus on methods that learn representations to produce optimal value functions for any given reward function. Rep- resentation learning through unsupervised self-supervised RI has been discussed for both pre-training (Nair et al.; Ma
2

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
et al., 2023) and training as auxiliary objectives (Agarwal et al., 2021; Schwarzer al., 2021a). While using auxiliary objectives for representation learning does accelerate policy learning for downstream tasks, the policy learning begins from scratch for a new task. Pre-training methods like (Ma et al., 2023; Nair et al.) use self-supervised learning tech- niques from computer vision like masked auto-encoding to learn representations that can be used directly for down- stream tasks. These methods use large-scale datasets (Grau- man et al., 2022) to learn representations but these are fitted around the policies used for collecting data. These represen- tations do not represent any possible behavior nor are trained to represent Q functions for any reward functions. Several prior works aim to discover intents or skills using a diversity objective. These methods use the fact that the latents or skills should define the output state-visitation distributions thus diversity can be ensured by maximizing mutual infor- mation ( (Warde-Farley et al., 2019; Eysenbach et al., 2019; Achiam et al., 2018; Eysenbach et al.. 2022) or minimizing Wasserstein distance (Park et al., 2024b) between the la- tents and corresponding visitaiisi distributions. PSM differs from these works and takes a step towards learning representations optimal for predicting value functions as well as a zero-shot optimal policy for any reward.
the optimal policy derived using Q function maximization which can lead to overestimation issues and instability dur- ing training as a result of Bellman optimality backups.
3. Preliminaries
In this section we introduce some preliminaries and define terminologies that will be used in later sections. We begin with some MDP fundamentals and RL preliminaries fol- lowed by a discussion on affine spaces which form the basis for our representation learning paradigm.
3.1. Markoy Decision Processes
A Markov Decision Process is defined as tuple (S, A, P.r, 7,p4) where S is the state space, A is the ac- tion space, P x A - A(S) is the transition proba- bility (A(.) denotes a probability distribution over a set), 7 E [0, 1) is the discount factor, μs is the distribution over initial states and r : S x A + R is the reward function. The task is specified using the reward function r and the initial state distribution μ.. The goal for the RL agent is to learn a policy Te : - A that maximizes the expected return ((T)) = -E F.r 1E: ag)]-
In this work, we consider a task-free MDP which does not provide the reward function or the initial state distribution. Hence, task-free or reward-free MDP simply the tuple (S, A, P,7). A task-free MDP essentially only captures the underlying environment dynamics and can have infi- nite downstream tasks specified through different reward functions.
Methods that linearize RL quantities: Learning basis vectors has been leveraged in RL to allow for transfer to new tasks. Successor features (Barreto et al., 2017) represents rewards as a linear combination of transition features and subsequently the Q-functions are linear in successor features. Several methods have extended successor features (Lehnert & Littman, 2020; Hoang et al., 2021; Alegre et al., 2022; Reinke& & Alameda-Pinedss 2021) to learn better policies in more complex domains.
The state-action visitation distribution, d' (s, a) is defined as the normalized probability of being in a state 8 and tak- ing an action if the agent follows the policy π from a state sampled from the initial state distribution. Concretely, dl (s,a) = (1 7) E 7 p'P(ss 8,0, =a. A more general quantity, successor measure, M a, st, aa), is defined as the probability of being in state st and taking ac- tion at when starting from the state-action pair s, a and fol- lowing the policy æ. Mathematically, M (s, a, a,,) = (1 - 7) Ei=o I'I = st,0c dr a+ |s0 = 8, 50 a). The state-action visitation distribution can be written as di (T(ss) a) E _so-n(()) | [M ag, s, a)}.
Spectral methods like Proto Value Functions (PVFs) (Ma- hadevan, 2005; Mahadevan & Maggioni, 2007) instead rep- resent the value functions as a linear combination of basis vectors. It uses the eigenvectors of the random walk op- erator (graph Laplacian) as the basis vectors. Adversarial Value Functions (Bellemare et al., 2019) and Proto Value Networks (Farebrother et al., 2023) have attempted to scale up this idea in different ways. However, deriving these eigenvectors from a Laplacian is not scalable to larger state spaces. (Wu et al., 2019) recently presented an approximate scalable objective, but the Laplacian is still dependent on random policy which usually makes it incapable of repre- senting all behaviors or Q functions.
Both these quantities, state-action visitation distribution and successor measure, follow the Bellman Flow equations:
dds(s, = (1 (1- +))))+++++++
Σ O'EA
((s)s', (I)
Similar to our work, Forward Backward (FB) Representa- tions (Touati& & Ollivier, 2021a; Touati et al., 2023) use an
For successor measure, the initial state distribution changes
inductive bias on the successor measure to decompose it into forward and backward representation. Unlike FB,
our representations are linear on set of basis features. Ad- ditionally, FB ties the reward with the representation of
3

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
to an identity function
tributions, it is possible to reduce the policy optimization problem to a search in this compact representation space. We will show that state visitation distributions and successor measures form an affine set and thus can be represented as E: 2: + b, where Φ; are basis functions, wx are "coor- dinates" or weights to linearly combine the basis functions, and b is a bias term. First, we build up the formal intuition for this statement and later we will use a toy example to show how these representations can make policy search easier.
Σ
Md (s, a, st ,at)= = (1 - aa=
Σ P(st a'))M"
(2)
E'ES
The RL objective has a well studied linear programming interpretation (Manne, 1960). Given any task reward func- tion r, the RL objective can be rewritten in the form of constrained linear program:
Σ
d(s, a))(s,)); s.1. d(s,a) 2 0 V8,a, a,
max
The first constraint in Equation 3 is the Bellman Flow equa- tion. We begin with Lemma 4.1 showing that state visitation distributions that satisfy the Bellman Flow form affine sets.
8,0
s.t. Σ d(ssa) = (1))
Theorem 4.1. All possible state-action visitation distribu- tions in an MDP form an affine set.
P(s|d (')d''a)dd(s'aa) o' E'
Σ
(3)
While Theorem 4.1 shows that any state-action visitation distribution in an MDP can be written using linear combi- nation of basis and bias terms, it still depend on the initial state distribution. Moreover, as shown in Equation 1, com- puting the state-action visitation distribution requires a sum- mation over all states and actions in the MDP which is not always possible. Successor measures are more general than state-action visitation distributions as they encode the visita- tion of the policy conditioned on a starting state-action pair. Using similar techniques, we show that successor measures also form affine sets.
and the unique policy corresponding to visitation d is d(s.0)
obtained by T[als) = Σ. d(s,aj The Q function can then be defined using successor measure as Q* a) = Est.a M a, s+, ++),st) at or Q" = Mr. Ob- taining the optimal policies requires maximizing the Q func- tion which requires solving arg max, Mr..
3.2. Affine Spaces
Let V be vector space and b be vector. An affine set is defined as A = b + V = {æ:|æ= = b + u,u E V}. Any vector in a vector space can be written as a linear combi- nation of basis vectors, i.e., v = Σ a;t; where n is the dimensionality of the vector space. This property implics that any element of an affine space can be expressed as b bbbb+ Σi Ogyy; Given system of linear equations Ar = c, with A being an X matrix (m < n) and * O, the solution x forms an affine set. Hence, there exists alphas a such that x= = b+ Σ; aqZ, The vectors {a;} } form the basis set of the null space or kernel of A. The values (os) form the affine coordinates of x for the basis {a,]. Hence, for a given system with known {a;} } and b, any solution can be represented using only the affine coordinates (ax).
Corollary 4.2. Any successor measure, M in an MDP forms an affine set and so can be represented as Et dewl + where Φ; and b are independent of the policy a and d is the dimension of the affine space.
Following Corollary 4.2, for any w. the function Σ S.w7 + b will be a solution of Equation 2. Hence, given &b (di stacked together) and b. we do not need the second con- straint on the linear program (in Equation 3) anymore. The other constraint: @.w; + b 2 still remains which w needs to satisfy. We discuss ways manage this constraint in Section 5.3. The linear program given a reward function now becomes,
max E + b)r]
We first explain the thcoretical foundations of our method in Section 4 and derive a practical algorithm following the theory in Section 5
t0
(4)
8.l. tu + +b> 0 Vs,a.
In fact, any visitation distribution that is a policy- independent linear transformation of M such as state visi- tation distribution or future state- visttviistttion distribution, can be represented in the same way as shown in Corollary 4.3.
4. The Basis Set for All Solutions of RL
In this section, we introduce the theoretical results that form the foundation for our representation learning approach. The proof for all the theoretical results can be found in Appendix 8. The goal is to learn policy-independent representations that can represent any valid visitation distribution in the environment (i.e. satisfy the Bellman Flow constraint in Equation 3). With a compact way to represent these dis-
Corollary 4.3. Any quantity that is policy-independent linear transformation of M can be written as a linear combination of
basis and bias terms.
Extension to Continuous Spaces: In continuous spaces, the basis matrices & and bias b become functions O 5xAx S + dd and b : S x A x S - R. The linear equation with matrix operations becomes a linear equation with functional
4

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
W;
dent basis functions, Φ0 to represent value functions as linear span, V - Some prior works (Dadashi et al., 2019) have already argued that value functions do not form convex polytopes. We show through Theorem 4.4 that for identical dimensionalities, the span of value functions using basis functions represent a smaller class of value func- tions than the set of value functions that can be represented using the span of the successor measure.
do
(a)
Theorem 4.4. Given a d-dimensional basis B : IR" + Rd define span ) as the span of all linear combinations of ba- sis B. Further define spana [Br] as the spun of inner prod- ucts of all linear combinations of basis B and all possible reward functions r. Let span (T"1) } denote the space of the value functions spanned by 52 while {spant 4 }}r} denotes the space of value functions using the successor measures spanned by 4. For the same dimensionality of task (policy or reward) independent basis, span($"11 S {span ( }r} for some 6.
(b)
Figure 2. (left) A Toy MDP with 2 states and 2 actions to depict how the linear program of RL is reduced using precomputation. (right) The corresponding, simplex for w assuming the initial state distribution is μ = (1,0)*.
transformations, and any sum over states is replaced with expectation under the data distribution.
Toy Example: Let's consider simple 2 state MDP (as shown in Figure 2a) to depict how the precomputation and inference will take place. Consider the state-action visitation distribution as in Equation 1. For this simple MDP, the T and b can be computed using simple alge- braic manipulations. For a given initial state-vise viste-viitatiinn distribution, fu and 7,, the state-action visitation distribu- tion d' (d(so, oo), d(s1, ao), d(s0, 01), d(s),a)))² can be written as,
Approaches such as Forward Backward Representations (Touati & Ollivier, 2021a) have also been based on repre- senting successor measures but they force a latent variable 2 representing the policy to be a function of the reward for which the policy is optimal. The linear one-to-one mapping between reward functions and corresponding policies are incorrect there can be many rewards leading to the same optimal policy as well as many optimal policies correspond- ing to single reward function. In addition, the forward map that they propose is a function of this latent z. We, on the other hand, propose a representation that is truly independent of the policy or the reward.
pppo)) 1+7 p|a H7pp)
d = W1
+ W2
+
14y
(5)
1
0
0
The derivation for these basis vectors and the bias vector can be found in Appendix A.6. Equation 5 represents any vector that is solution of Equation 1 for the simple MDP. Any state-action visitation distribution possible in the MDP can now be represented using only w (w,, w²)"". The only constraint in the linear program of Equation 4 is + b 2 O. Looking closely, this constraint gives rise to four inequalities in w and the linear program reduces to,
5. Method
In this section, we start by introducing the core practical algorithm for representation learning inspired by the theory discussed in Section 4 for obtaining & and b. We then discuss the inference step, i.e., obtaining w for a given reward function.
yW1 W2
W1 7/222
Tr
max w1,w2
w1,
1 1
+
5.1. Learning + and b
s.t.
W1 + yW2 S p(50) + 7H($1)
For a given policy n, its successor measure under our frame- work is denoted by M ow + with w the only object depending on policy. Given an offline dataset with density p. follow prior works (Touati & Ollivier, 202 1a; Blier et al., 2021b) and model densities m.* = M² /p learned with the following objective:
ywn + w2 S p(s1) + 71000)
Wi 2 0, w2
The inequalities in w give rise to a simplex as shown in Fig- ure 2b. For any specific instantiation of μ and r, the optimal policy can be easily found. For instance, if μ (1.0)" and the reward function, r (1, 0, 1, 0)T, the optimal will be obtained at the vertex (w) = 1, W2 = 0) and the correspond- ing state-action visitation distribution is d = (0, 0, T) As shown for the toy MDP, the successor measures form a simplex as discussed in (Eysenbach et al., 2022). Spectral Methods following Proto Value Functions (Mahadevan & Maggioni, 2007) have instead tried to learn policy inedpnn-
I. I.7 w (1 B), 0. 0,b.0" s,)))
E B. .o p b.b. (8, sta,sttaa) ,at) 2
+
ym h.b00 (5) ((')), aat))) (7)
The above objective only requires samples (s, a, 8') from the reward-free dataset and a random state-action pair
5

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
(also sampled from the same data) compute [(x).
that is more stable than the one derived by maximization (Farebrother et al., 2023).
A & and b that allows for minimizing the (TTr for all E I1 forms a solution to our representation learning problem. But how do we go about learning such & and 6? A naive way to implement learning & and b is via a bi-level optimization. We sample policies from the policy space of II, for each policy we learn ut that optimizes the policy evaluation loss (Eq 7) and take gradient update w.r.t & and b. In general, the objective can be optimized by any two-player game solving strategies with [ b] as the first player and "" as the second player. Instead, in the next section, we present an approach to simplify learning representations to a single-player game.
5.3. Fast Inference on Downstream Tasks
After obtaining & and b via the pretraining step, the only parameter to compute for obtaining the optimal Q function for downstream task in the MDP is w. As discussed earlier, Q' = max (pw + b)r but simply maximizing this objective will not yield a Q function. The linear program still has a constraint of Qw + 6 2 O, Vs, a. We solve the constrained linear program by constructing the Lagrangian dual using Lagrange multipliers X(s, a). The dual problem is shown in Equation 10. Here, write the corresponding loss for the constraint as min + b, O).
5.2. Simplifying Optimization via a Discrete Codebook of Policies
max min - -wur- MMs, a) min (ww + +,,0), (10) A20
8,0
Learning a new w for each specific sampled policy æ does not leverage precomputations and requires retraining from scratch. We propose parameterizing to be conditional on policy, which allows leveraging generalization between poli- cies that induce similar visitation and as we show, will allow us to simplify the two player game into a single player opti- mization. In general, policies are high-dimensional objects and compressing them can result in additional overhead. Compression by parameterizing policies with a latent vari- able z another alternative but presents the challenge of covering the space of all possible policies by sampling Instead, we propose using discrete codebook of policies as a way to simulate uniform sampling of all possible policies with support in the offline dataset.
Once w² is obtained, the corresponding M and Q" can be easily computed. The policy can be obtained as ' arg max, (* a) for discrete action spaces and via DDPG style policy learning for continuous action spaces.
6. Connections to Successor Features
In this section, we uncover the theoretical connections be- tween PSM and successor features. Successor Features (Bar- reto et al., 2017) (v) (s, a)) are defined as the discounted sum of state features p(s), (s.a)
E, [E, 7 (s))) These state features can be used to span reward functions as = pz. Using this construction, the Q function is linear in z as Q a) = v/t (s, a)z. We can establish a simple relation between M and ei²", 1ht (s, a) = Ssr M a, (()) (s(s')d''. This connection shows that, like successor measures, suc- cessor features can also be represented using a similar basis.
Discrete Codebook of Policies: Denote z as a compact representation of policies. We propose to represent as random sampling seed that will generate a determinis- tic policy from the set of supported policies as Equation 8. In other words, z will be a random integer (z S 2h (represented using h bits) that will represent a deterministic policy.
Theorem 6.1. Successor Features (8, a) belong to an affine set and can be represented using a linear combination of basis functions and a bias.
[[aas, z) = Uniform Sample(seed = z + hash())). (8)
Interestingly, instead of learning the basis of successor mea- sures, we show below that PSM can also be used to learn the basis of successor features. While traditional successor feature-based methods assume that the state features yo are provided, PSM can be used to jointly learn the successor feature and the state feature. We begin by introducing the following Lemma 6.2 from (Touati et al., 2023) which con- nects an a specific decomposition for successor measures to the ability of jointly learning state features and successor representations,
The above sampling strategy defines unique mapping from a seed to a policy. If the seed generator is unbiased, the approach provably samples from among all possible deter- ministic policies uniformly. Now, with policy и= and wz parameterized as a function of z we derive the following single-player reduction to learn 4 b, w jointly.
PSM-ob jective: arg min [L* (L, b, w(2))). (9) (x))
The disentanglement between the policies sampled during training and the corresponding reward makes Equation 9 more stable than similar looking objectives like FB (Touati & Ollivier, 2021b). Equation 9 is absent of any policy learn- ing via maximization of a moving reward function. In partic- ular, PSM uses a relatively mild form of off off-policy learning
Lemma 6.2. (Theorem 13 of (Touali et al., 2023)) For an offline dataset with density p. if the successor measure is represented as M= a, s*) = v5t a)p((s pp(st) then vp is the successor feature vt (s, for state feature (p(s)" (Bpp(ppT)))-)
6

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
According to Lemma 6.2, if (s, 4ht (s),)) )p(ss)) then the corresponding suc- cessor feature is 1 and the state feature is PSM represents successor measures as M 0, 8+) = o(s, a, )w= ((**) (for simplicity, combining the bias within the basis without loss of generality). can be shown that if the basis learned for successor measure using PSM. ( a,s.a,s, a, s*) is represented ()) forms
Baselines:
We compare against methods that commonly used and are the state of the art in spanning the space of reward functions: Laplacian features (Wu et al., 2018), Forward Baaddd-ra (Touati et al., 2023) and HILP (Park et al., 2024a). Laplacian features learn features of a state by considering eigenvectors of a graph Laplacian induced by a random walk. These features ()(s) E Rd obtained for each state are used to define a reward function conditioned on (s) 2 where is sampled uniformly
((s)""((((p"))
as decomposition (s(s,a) p(st), the basis for successor features for the state features p(s)) (Ep((())))) Formally, we present the following theorem,
reward r(s; from a unit d-dimensional sphere. For each an optimal policy is pretrained from the dataset on the induced reward function. During inference the corresponding z for a given reward function is obtained as a solution to the following linear regression: min; [.
Theorem 6.3. For the PSM representation M as, o(s, st +tww" and Φ(8, a, s+) ++) (s, aa"aa)""""''")))))p))) cessor feature zit (s. a) = (5 a) fn
r(s)) Similar to Laplacian features, we consider a couple other features in the SF framework, one that uses one-step forward dynamics predictability (method named "FDM") and one that uses SVD decomposition of successor representation (method named "SVD"). Forward-backwaid (FB) learns both the optimal policy and state features jointly for all reward that are in the linear span of state-fcatures. FB methods typically assume a goal-conditioned prior during pretraining which typically helps in learning policies that reach various states in the dataset. HILP (Park et al., 2024a) makes two changes to FB: a) Reduces the tasks to be goal reaching and b) Uses more performant offline RL method, IQL (Kostrikov et al., 2021) ) to learn features. We provide detailed experimental setup and hyperparameters in Appendix B.3.
((*)), ), the suc- a) for the state feature
(Ep((("))))")))
Thus, successor features can be obtained by enforcing a particular inductive bias to decompose é in PSM. For rewards linear in state features (r(s) = (s)()) =) for some weights z), the Q-functions remain linear given by Q*(s,a) a) dw (s, a) (+))==. A natural question to ask is, with this decomposition, do we lose the express- ibility of PSM compared to the methods that compute ba- sis spanning value functions, thus contradicting Theorem 4.4? The answer is negative, since (1) even though the value function seems to be linear combination of some ba- sis with weights w",, these weights are not tied to 2 or the reward. The relationship between the optimal weights w= and z defining the reward function is not necessarily lin- ear as the prior works assume, and (2) the decomposition Ø(s, a, s+) = (s, ())(*)))) reduces the representation capacity of the basis. While prior works are only able to recover features pertaining to this reduced representation capacity, PSM does not assume this decomposition and can learn larger representation space. Additionally, a number of SF methods assume access these features or learn them using auxiliary objective, assuming that the obtained fea- tures would be sufficient to span successor measures. On the other hand, PSM is able to extract the features relevant for learning M for all π.
7.1. Zero shot Value function and Optimal Policy prediction
7. Experimental Study
Loptoos
Our experiments evaluate how PSM can be used to encapsu- late a task-free MDP into a representation that will enable zero-shot inference on any downstream task. the ex- periments we investigate a) the quality of value functions learned by PSM (Section 7.2, b) the zero-shot performance of PSM in contrast other baselines discrete tasks (Sec- tion 7.1 and Appendix C.2), c) the ability to learn general goal-reaching skills arising from the PSM objective on a robot manipulation task (Section7. 2)and finally d) Suitabil- ity of learned PSM representations for enabling zero-shot RL continuous state-action space tasks. (Section 7.3)
(a) GGidworld
(b) Four-room
Figure 3. Qualitative results on a gridworld and four-room:
denotes the goal sampled for every episode. The black regions are the boundariesoobstacles. The agent needs to navigate across the grid and through the small opening (in case of four-room) to reach the goal. We visualize optimal Q-functions inferred at test time for the given goal in the image. The arrows denote the optimal policy. (Top row) Results for PSM, (Middle Row) Results for FB, (Bottom row) Results for Laplacian Eigenfunctions.

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
Task
Laplace
FDM
SVD
FB
HILP
PSM
Walker
Stand
243,70
151.40
834,42 + 39.36
822.52 + 99.25
902.63-4: 38.94
6U7.07 + 165,28
872.61 + 38.81
Run
63.65
31,02
269.27 1 27.52
332.06
1: 12.87
392.76
1: 31.29
107.84
34.24
351.50
19.46
Walk
190.53
168.45
647.76 +: 180.99
852.06
13:67
877.10 1 81.05
399.67
=39.31
891.44
46.81
Flip
48.73 = 17.06
441.09 + 66.48
444.43
+ 106.56
206.22 = 162.27
277.95 + 59.63
640,75 + 31.88
Averagel*)
136.65
548.14
612.76
594.67
348.13
689.07
Cheetah
Run
96.32 = 35.69
161.90 63.51
171.55 + 15.05
257.59 + 58.51
68.22 4:47.08
244.38 80.00
Ron Backward
106.38 +: 29.4
280.16 + 17.76
199.79
= 28.05
307.07 + 14.91
37.99 4:25,
296.44 + 20.14
Walk
d: 56.08
+: 177.14
672.03
318 ::
984,21 0.49
Walk Backward
654,29 219.81
945.91 1: 63.13
858.33
d: 42.89
980.76 +: 2.32
349.61 1: 236.29
979.01 = 7.73
Average(*)
316.53
494.46
475.42
586.31
193.53
626.01
Quadruped
854.50 1: 41.47
940.01
894.30 35.79
740.05
107.15
409.54
97.59
842.86
Run
412.98 +: 54.03
434.08 + 21.85
433.66 + 21.89
386.67 + 32.53
205.44 +: 47.89
431,77 + 44,69
Walk
494.56 + 62.49
459.43 32.44
457.74 64.93
566.57 + 53.22
218.54 86.67
606.97=73.67
Jump
642.84 114.15
709.07
+ 71.99
687.75
40.59
581.28 = 107.38
325.51
=93.06
596,37 #:4423
Average(*)
601.22
635.65
618.36
568.64
289.75
618.74
Pointmass
Topp Left
713.46 +: 58.90
891.56 49.16
853.79
35.49
897.83 +: 35.79
944.46 12.94
831.43
69.51
Top Right
581.14
214.79
379.75
84.38
442.92 1: 123.84
274.95
197.90
96.04 1: 166.34
730.27
58.10
Bottom Left
689.05 +: 37.08
563.53 + 213.60
$17.19 + 142.36
517.23 + 302.63
192. 34 +: 177.48
451.38 +: 73.46
Bottom Right
21.29 42.54
10.96 + 13.51
57.74
80.87
19.37=33.54
0.17 = 0.29
43.29 + 38.40
Average(*)
501.23
461.45
467.91
427.34
308.25
514.09
Table Table shows comparison (over 5 seeds) of zero-shot RL performance between different methods with representation size of d= = 128. PSM demonstrates marked improvement over prior methods. (*) denotes statistically significant through Hann-Whitney U Test with level 0.05.
In this section, we consider goal-conditioned rewards on discrete gridworld and the classic four-room environments.
7.2. Learning zero-shot policies for manipulation
We consider the Fetch-Reach environment (Figure 4) with continuous states and discrete actions (Touati & Ol- livier, 2021a). A dataset of size 1M is constructed using DQN+RND. FB, Laplacian and PSM all use this dataset to learn pretrained objects that can be used for zero-shot RL.
Task Setup: Both environments have discrete state and action spaces. The action space consists of five actions: {up, right, down, left, stay]. We collect transitions in the environment by uniformly spawning the agent and taking random-uniform action.This allows us to form our of- fline reward-free free dataset will full coverage to train & and b. During inference, we sample goal and infer the opti- mal Q function on the goal. Since the reward function is given by r(s) = -m the inference looks like Q(s, a) = max ((s, a. g)w s.t. D(s, a, s') + b(s, s's) > O Vs,a, s'. Figure shows the Q function and the cor- responding optimal policy (when executed from a fixed start state) on the gridworld and the four-room environment. As illustrated clearly, for both the environments, the optimal Q function and policy can be obtained zero-shot for any given goal-conditioned downstream task. The error rate of each of the method on these tasks are presented in Appendix C.2
We observe that PSM outperforms baselines FB and Lapla- cian in its ability to learn a zero-shot policy. One key obser- vation is that PSM learning is stable whereas FB exhibits drop in performance, likely due to the use of Bellman optimality backups resulting in overestimation bias during training. Laplacian's 's capacity to output zero-shot policies is far exceeded by PSM because Laplacian methods construct the graph Laplacian for random policies and may not be able represent optimal value functions for all rewards.
7.3. Learning Zero-shot Policies for Continuous Control
We use the ExoRL suite (Yarats al., 2022) for obtaining exploratory datasets collected by running RND (Burda et al., 2019). PSM objective in Equation 9 directly enables learn- ing the basis for successor measures. We decompose the basis representation d(s, a, 8+) to ths a) so(+))))))))))))))))))) )) as dis- cussed in detail in Section 6. PSM thus ensures that yos*t can be used to construct basic features to span any reward function. Note that this is not limiting assumption, as the features can be arbitrarily non-linear in states. In these experiments, we compare the ability of PSM to obtain these representations as compared to prior zero-shot RL methods. Experimental details can be found in Appendix B.3.
Comparison to baselines: We can draw a couple of con- clusions from the visualization of the Q functions inferred by the different methods. First, the Q function learnt by PSM is more sharply concentrated on optimal state-action pairs compared to the two baselines. Both baselines have more uniform value estimates, leaving only a minor differ- ential over state values. Secondly, the baselines produce far more incorrect optimal actions (represented by the green arrows) compared to PSM.
Table compares PSM's zero performance in contin-
8

---

Proto Successor Measure: Representing the Behavior Space of RL Agent
outperforming prior baselines.
Limitations and Future Work: PSM shows that any MDP can be compressed to a representation space that allows zero-shot RL, but it remains unclear as to what the size of the representation space should be. A large representational dimension can lead to increased compute requirements and training time with possible chance of overfitting, and a small representation dimension can fail capture nuances about environments that have non-smooth environmental dy- namics. An interesting future direction would be to study the impact of dataset coverage on zero-shot RL performance.
FetchReech environment
1.-
0.8
Acknowledgements
Success rate
0.6
The authors would like to thank Scott Nickum, Ahmed Touati, Ishan Durugkar and Ben Eysenbach for inspiring discussions during this work. We thank Caleb Chuck, Max Rudolph, and members of MIDI Lab for their feedback on this work. SA, HS, and AZ are supported by NSF 234065 I and ARO W911NF-24-1-019.. This work has in part taken place in the Learning Agents Research Group (LARG) at the Artificial Intelligence Laboratory, The University of Texas at Austin. LARG research is supported in part by the Na- tional Science Foundation (FAIN-201984) NRT-2125858) the Office of Naval Research (N00014-15 1800001-18-22444 18-2243), Army Re- search Office DARPA (Cooperative Agreement HR001 12520004 on Ad Hoe Teamwork), Lockheed Martin, and Good Systems, a re- search grand challenge at the University of Texas at Austin. The views and conclusions contained in this document those of the authors alone. Peter Stone serves the Exec- utive Director of Sony AI America and receives financial compensation for this work. The terms of this arrangement have been reviewed and approved by the University of Texas at Austin accordance with its policy on objectivity in re- search.
0.9
0.2
psm
leplace
0.0
20000
20000
SCOO0
Update steps
Figure 4. Quantitative results on FetchReach: The success rates (averaged over 3 seeds) are plotted (along with the standard devia- tion as shaded) with respect to the training updates for PSM, FB and Laplacian. PSM quickly reaches optimal performance while FB shows instability in maintaining its optimality. Laplacian is far from the optimal performance.
W91 W911NFF-1-2-1 17-2-1188-111111 81),
uous state-action spaces to representative methods Lapla- cian, FB, and HILP. We note that to make the comparisons fair, we use the same representation dimension of d 128, the same discount factor, and the same inference and pol- icy extraction across environments for particular method. Overall, PSM performs consistently better or is competitive to baselines across the environments. Ablations studying effect of latent dimensionality can be found in Appendix C.
Impact Statement
8. Conclusion
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences our work, none which we feel must be specifically highlighted here.
This paper introduces Proto Successor Measures (PSM), zero-shot RL method that compresses any MDP to allow for optimal policy inference for any reward function with- out additional environmental interactions. This framework marks a step in the direction of moving away from com- mon ideology in RL to solve single tasks optimally, and rather pretraining reward-free agents that are able to solve an infinite number of tasks. PSM based on the principle that successor measures are solutions to affine set and proposes an efficient and mathematically grounded algo- rithm to extract the basis for the affine set. Our empirical results show that PSM can produce the optimal Q function and the optimal policy for a number of goal-conditioned as well as seeerf-s tasks number of environments
References
Achiam, J., Edwards, H., Amodei, D., and Abbecl, P. Variational option discovery algorithms. CORR, ha//1807.10299 2018. URL httt:://axxvv. abs Abs//807.10253
Agarwal, S., Courville, A., and Agarwal, R. Behavior pre- dictive representations for generalization in reinforce- ment learning. In Deep RI. Workshop NeurIPS 2021,

---

Proto Successor Measure: Representing the Behavior Space of RL Agent
2021. URL httppp:://peeenereee..eeeeteeeee id=b5P.aax56Jx.
Eysenbach, B., Salakhutdinov, R., and Levine, S. The infor- mation geometry of unsupervised reinforcement learning. In International Conference Learning Representations, 2022. URL https:/ https:://ppereeeeeeeee wpwppeveeeeeeeeeeeeeeeeeeeeeeeee net forum? i id=3w02UX0voF.
Agarwal, S., Durugkar, I., Stone, and Zhang, A. I-policy gradients: A general framework for condicondi rl using f-divergences. Advances in Neural Information Processing Systems, 36, 2024.
Farebrother, J., Greaves, J., Agarwal, R., Lan, C. L., Goroshin, R., Castro, P. S., and Bellemare, M. G. Proto- value networks: Scaling representation learning with auxiliary tasks. The Eleventh International Confer- ence on Learning Representations, 2023. URL https: / //openeeveww net http:/rrmmmooooatttttt-----
Alegre, N., Bazzan, A., and Da Silva, B. C. Optimistic lin- car support and successor features as a basis for optimal policy transfer. In International conference on machine learning, pp. 394-413, PMLR, 2022.
Barreto, A., Dabney, W., Munos, R., Hunt, J. J., Schaul, T., van Hasselt, H. P., and Silver, D. Successor features for transfer in reinforcement learning. Advances in neural information processing systems, 30, 2017.
Fawzi, A., Balog, M.. Huang. A., Hubert, T., Romera- Paredes, B., Barekatain, M., Novikov, A., R Ruiz, F. J., Schrittwieser, J., Swirszez, G., et al. Discovering faster matrix multiplication algorithms with reinforce- ment learning. Nature, 610(7930):47-5,,22 2022.
Bellemare, M., Dabney, W., Dadashi, R., Ali Taiga, A., Castro, P. S., Le Roux, N., Schuurmans, D., Lattimore, T., and Lyle, C. A geometric perspective on optimal representations for reinforcement learning. Advances in neural information processing systems, 32, 2019.
Ghosh, D., Bhateja, C. A., and Levine, S. Reinforcement learning from passive data via latent intentions. Inter- national Conference on Machine Learning, pp. 11321- 11339. PMLR, 2023.
Blier, L., Tallec, C., and Ollivier, Y. Learning successor states and goal-dependent values: A mathematical view- point. arXiv preprint arXiv:2 21100111111.,... 2021a.
Grauman, K., Westbury, A., Byrne, E., Chavis, Z., Furnari, A., Girdhar, R., Hamburger, J., Jiang, H., Liu, M.. Liu, X., et al. Ego4d: Around the world in 3,000 hours of egocentric video. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 18995-19012, 2022.
Blier, L., Tallec, C., and Ollivier, Y. Learning succes- sor states and goal-dependent values: A mathemati- cal viewpoint. CoRR, abs/2101.07123, 2021b. URL. https: pettps://arxien /tt///axiivv...//////eee/eaaaavm https///axxvvv../aaae/ee
Hafner, D., Lillicrap, T., Ba. J., and Norouzi, M. Dream control: Learning behaviors by latent imagination. In International Conference Learning Representations, 2020. URL https:/
Burda, Y., Edwards, H., Storkey, A., and Klimov, O. Exploration by random network distillation. In In- ternational Conference on Learning Representations, 2019. URL https://ppen.rreiee wwwpeeeiiiw.eeeteeeeeeeeeeeeeeee net//frunt id=H11JJnR5W
httpp://opennevvai hwwpe/wveew/eeeeceoe/eeeeeeeeeep net net/oouum forum? id=s110TC4tDS.
Hoang, C., Sohn, S., Choi, J., Carvalho, W. T., and Lee, H. Successor feature landmarks for long-horizon goal- conditioned reinforcement learning. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, W. (eds.), Advances Neural Information Processing Systems, 2021. URL https: //wpnnevieiwweeeeeeoeeeeee net/ forum? id=rD6u12FTbf.
Dadashi, R., Taiga, A. A., Le Roux, N., Schuurmans, D., and Bellemare, M. G. The value function polytope in reinforcement learning. In International Conference Machine Learning, pp. 1486-1495. PMLR, 2019.
Dayan, P. Improving generalization for temporal difference learning: The successor representation. Neural computa- tion, 5(4):613-624, 1993.
Janner, M., Fu, J., Zhang, M., and Levine, S. When trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
Denardo, E. V. On linear programming in a markov decision problem. Management Science, (6(5):281-288) 1970.
Koren, Y. On spectral graph drawing. In International Computing and Combinatorics Conference, pp. 496-508. Springer, 2003.
Eysenbach, B., Gupta, A., Ibarz, J., and Levine, S. Diversity is all you need: Learning skills without reward function. In International Conference on Learning Representatonss, 2019. URL https: : //pp/reeeeveeveee htt.rviwwwweeeeeeeeeeeeeeeee forum? dd= 63 JRFm.
Kostrikov, I., Nair, A., and Levine, S. Offline reinforce- ment learning with implicit q-learning. arXiv preprint arXiv:2 10,06169, 2021.
10

---

Proto Successor Measure: Representing the Behavior Space of RL Agent
Lehnert, L. and Littman, M. L. Successor features combine elements of model-free and model-based re- inforcement learning. Journal of Machine Learning Research, 21(196):1-53, 2020. URL http:// jmlr. org/paperse / /v/1/12111///11// 111222119000000 9-060.. html.
Twelfth International Conference on Learning Represen- lations, 2024b. URL /ttppr/vieeeeeeeeeee.e/// net/ forum?dd-c5pwOSS00
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. Learning transferable visual models from natural language supervision. CoRR, abs/2103 000200, 2021. URL https: /ttp://aaxvivoooeeeeo a abs 20../213.000 000200..
Ma, Y. J., Sodhani, S., Jayaraman, D., Bastani, O., Kumar, V., and Zhang, A. VIP: Towards universal visual reward and representation via value-implicit pre-training. In The Eleventh International Conference on Learning Represen- tations, 2023. URL https: http/ppeeeevveeeeeeeeeeeeeeeeeee net./ forum? dddoddot7.........
Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Rad- ford, A., Chen, M., and Sutskever, Zero-shot text-to- image generation. CoRR, abs/2102..2099. 2021. URL https: httpp:://azaav..rrg/aaa.. http://aas///1..22222 12092
Machado, M. C., Bellemare, M. G., and Bowling, M. A laplacian framework for option discovery in reinforce- ment learning. In International Conference on Machine Learning, pp. 2295-2304. PMLR, 2017.
Reinke, C. and Alameda- X. Xi-learning: Successor feature transfer learning for general reward functions. CoRR, abs/21 2021. URL htthhttp://aae
Machado, M. C.,, Rosenbaum, C.,, Guo, X., Liu, M., Tesauro, G., and Campbell, M. Eigenoption discovery through the deep successor representation. In International Confer- ence on Learning Representations, 2018. URL https: //openreview wwt.rriwwwweeeeeeeeeeee.ee..eeee net//oruum id=Bk BZCAXR-
Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., and Bachman, P. Data-efficient reinforce- ment learning with self- pelf-predietive representations. In International Conference on Learning Representations, 2021a. URL https:/ wpw.reveeeeeeeeeeeeeeeeee..oe.wm net ntt.:/wuweee///rr forum? i id=uCQfPZwRau
forum?
Mahadevan, S. Proto-value functions:
Developmental re- inforcement learning. In Proceedings of the 22nd inter- national conference on Machine learning. pp. 553-560, 2005.
Schwarzer, M., Rajkumar, N., Noukhovitch, M., Anand, A., Charlin, L., Hjelm, R. D., Bachman, P., and Courville, A. C. Pretraining representations for data-efficient re- inforcement learning. Advances in Neural Information Processing Systems, 34: 12686-12699, 2021b.
Mahadevan, S. and Maggioni, M. Proto-value functions: A laplacian framework for learning representation and control markov decision processes. Journal of Machine Learning Research, 8(10), 2007.
Manne, A. S. Linear programming and sequential decisions. Management Science, 6(3):259-267, 1960.
Sermanet, P., Lynch, C., Hsu, J., and Levine, S. Time- contrastive networks: Self-supervised learning from multi-view observation. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pp. 14-15,201 2017.
Mnih, V., Kavukcuoglu, K., Silver, D., Graves, A., Antonoglou, I., Wierstra, D., and Riedmiller, M. A. Playing atari with deep reinforcement learning, CoRR, abs/131,.5600,, 2013. URL abs/1312.560
Sikchi, H., Chitnis, R., Touati, A., Geramifard, A., Zhang, A., and Niekum, S. Score models for offline goal- conditioned reinforcement learning. The Twelfth International Conference on Learning Representations, 2024a. URL https: /opp/reveeeeeeeeeeeeeeeeeeeeeeee net forum? i 1d=oX]nwQT.CTA.
Nachum, O. and Dai, B. Reinforcement learminngg via fenchel- rockafellar duality. arXiv preprint arXiv:2001.011 200600000,0186,, 2020.
Nair, S., Rajeswaran, A., Kumar, V., Finn, C., and Gupta, A. R3m: A universal visual representation for robot manipu- lation. In 6th Annual Conference Robot Learning.
Sikchi, H., Zheng, Q., Zhang, A., and Niekum, S. Dual RL: Unification and new methods for reinforcement and imitation learning. The Twelfih International Confer- ence on Learning Representattonss, 2024b. https: / //openeevew..r / picerrumiidsse 9Bu66rqv.
Park, S., Kreiman, T., and Levine, Foundation policies with hilbert representations. In Forty-first International Conference on Machine Learning, 2024a. URL https: //openreview.co net/ forumm Saw.rrummii-hhhi-eee Li http:/viwwinnwteeeeeeeeeeeee
Stooke, A., Lee, K., Abbeel, P., and Laskin, M. Decoupling representation learning from reinforcement learning. In International conference on machine learning, pp. 9870 9879. PMLR, 2021.
Park, S., Rybkin, O., and Levine, S. METRA: Scalable unsupervised RI. with metric-aware abstraction. In The

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., de Las Casas, D., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., Lillicrap, T. P., and Riedmiller, M. A. Deepmind control suite. CoRR, htts/1801.000900 2018. URL petttp://arxiee o abs / 801 00690.
Touati, A. and Ollivier, Y. Learning one representation optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23, 2021a.
Touati, A. and Ollivier, Y. Learning one representation optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23 2021b.
Touati, A., Rapin, J., and Ollivier, Does zero-shot reinforcement learning exist? In The Eleventh In- ternational Conference on Learning Representations, 2023. URL httpp:://ppereetteee wtwpeeeiewwneeeeeeeeeeeeeeeeeeee id=MYFap_CcCB
Warde-Farley, D., de Wiele, T. V., Kulkarni. T., Ionescu, C., Hansen, S., and Mnih, V. Unsupervised control through nonpparamttric discriminative rewards. In In- lernational Conference on Learning Representations, 2019. URL https: http//ppnrrveeeeeeeeeeereeeeeeee id=rleVMnA9K7.
Wu, Y., Tucker, G., and Nachum, O. The laplacian in rl: Learning representations with efficient approximations. 2018
arXiv preprint
Wu, Y., Tucker, G., and Nachum, O. The laplacian in RL: Learning representations with efficient approximations. In International Conference on Learning Representations, 2019. URL https: wtwpeerrreeew.eeeeeeeeeeeeeeeeee id=HJ1 NpoASQQ.
Wurman, P. R., Barrett, S., Kawamoto, K., MacGlashan, J., Subramanian, K., Walsh, T. J., Capobianco. R., Devlic, A., Eckert, F., Fuchs, F. et al. Outracing champion gran turismo drivers with deep reinforcement learning. Nature, 602(7896):23--288 2022.
Yarats, D., Brandfonbrener. D., Liu, H., Laskin, M., Abbeel, P., Lazaric, A., and Pinto, L. Don't 't change the algorithm, change the data: Exploratory data for offline reinforce- ment learning. arXiv preprint arXiv:2201. 13425, 2022.
Zahavy, T.. Schroecker, Y., Behbahani, F., Baumli, K., Flen- nerhag, S., Hou, S., and Singh, S. Discovering poli- cies with DOMiNO: Diversity optimization maintaining near optimality. The Eleventh International Confer- ence on Learning Representations, 2023. URL https: / //openreview.no net wormm/ormmmm///mm 1eae//ooormm111ekkkk2223333333
12

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
Appendix
A. Theoretical Results
In this section, we will present the proofs for the Theorems and Corollaries stated Section and 6.
A.1. Proof of Theorem 4.1
Theorem 4.1. All possible state-action visitation distributions in an MDP form an affine set.
Proof. Any state-action visitation distribution, (s, (s,a) must satsify the Bellman Flow equation:
Σ dd = (1-7))(s)) +
Σ P(s|s a (s',a'')))) (s',
(11)
s'.a'
This equation can be written in matrix notation as:
Σ dh = - (1 + 7P d=
(12)
Rearranging the terms,
(S 7P' ddt (1
(13)
where S is the matrix for E. of size |S] x [8|| with only |A| entries set corresponding the state denoted by the row. This equation is affine equation of the form Ax whose solution set forms affine set. Hence all state-visitation distributions d' form an affine set.
In the continuous spaces, the visitation distributions would be represented as functions: dt S x A + R rather than vectors in [0, 1|S The state-action visitation distribution d (8, a) will satisfy the following continuous Bellman Flow Equation,
J, d" (s, a)da = ()ms) + J
P(s]s', aa)dd (s', yd''d'''
(14)
This equation is the same as Equation 11 except, the vectors representing distributions are replaced by functions and the discrete operator Σ is replaced by f-
The Bellman Flow operator can be defined T that acts on dm as.
(s) - /f d= (s, a) -? / (Pssss, d')dd (8', d''ds'da da'.
(15)
From Equation 14, [s) = (1 (1 (s). The operator a linear operator, hence d' (s,a) forms an affine space.
A.2. Proof of Corollary 4.2
Corollary 4.2. Any successor measure, M in an MDP forms an affine set and so can be represented as Σ 0: + b where di and b are independent of the policy æ and d the dimension of the affine space.
Proof. Using Theorem 4.1, we have shown that state- acte-nction visitation distributions form affine sets. Similarly, successor measures, (s a, st aa) are solutions of the Bellman Flow equation:
(s, a, sttttttt++))))) = (1 1[s = st, a = a+1++ + Σ
P(s (s'ttss'''))'( s d')r(a+lst))
(16)
s'.''ESA
Taking summation over a on both sides gives us an equation very similar to Equation 11 and so can be written by rearranging as,
(S- - > = (1 - -
(17)
With similar arguments as in Lemma 4.1, M also forms an affine set.
Following the previous proof, continuous spaces, R becomes function : A + and the Bellman Flow equation transforms to,
M M as, st, s*.tt)*=++++++ (1 >)ppp(s =t +
d')r(aat)s'tt dd'
(s( da'. (18)
S
13

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
Integrating both sides over at, the Bellman Flow operator T
can be constructed that acts on M", [s,"][s,,,s7]] / (s, st aa "ddad -7 -ddaddaaaa
/ Pls (s, a')d
(19)
(s) = (1 ()p(s= st a aa=aa =at)
(20)
As T is linear operator, M belongs an affine set.
Any element of affine set of dimensionality d, can be written Σi o:w; b where (dx) are the basis and b is bias vector. The basis is given by the null space of the matrix operator (S PPT (T in case of continuous spaces). Since the operator (S P (and T' and the vector (1 1) st] (and function (1 7)p(( = st, ==== a+)) are independent of the policy, the basis & and the bias are also independent of the policy.
A.3. Proof of Theorem 4.4
Theorem 4.4. Given d-dimensional basis B IR + Rd, define [ as the span of all linear combinations of basis B. Further define span ( B as the span of inner products of all linear combinations of basis B and all possible reward functions r. Let span| denote the space of the value functions spanned by Φ- while {spant { denotes the space of value functions using the successor measures spanned by Φ. For the same dimensionality of task (policy or reward) independent basis, span (Te5) (s)an(D)r) for some +
Proof. We need to show that any element that belongs to the set spant [ also belongs to the set (span { }r}.
Any element belonging to the set (span((")) {}" is represented by,
(T(() = _ S7 87 (s).
3T can be written as R.
0r k; where ki = =.. r(s") for some r. This means, (s) = Σ k r(s')
Σ w (((s))))((s)
where w. = T (s( 8. (s) This implies, for every instance of V E span there exists some k.
instance in (span(0)) (r) for some t and r.
Lets see when an element in {span{ (rr belongs in span( We start from an element belonging to set {span { }r} as represented by,
(*(() = Σ i Σ ('sss''rrss))
If we assume a special (,(s, (s)) (T)
Σ To Σ (r)
Σ Tw En(s)r(s))osts
Σ 87 (s)
where = [w] E. (s))r(ss)) and (s) = (s). This implies for only for the special case 2,(s, = (s))) Value functions belong to the set spon ( (f and in general, this may not hold.
A.4. Proof of Theorem 6.1
Theorem 6.1. Successor Features (s, a) belong to an affine set and can be represented using a linear combination of basis functions and bias.
14

---

Proto Successor Measure: Representing the Behavior Space of RL Agent
Proof. Given basic state features, yp S Rl4, the successor feature is defined as, (s, En I 7 ((441)) I be correspondingly connected successor measures 10th (s, a) Eg M(s,a, s' (s((') (replace Ev with Sg for continuous domains). In Linear algebra notations, let M be a (8 x A) S dimensional matrix representing successor measure. Define P, as the S x d matrix containing for each state concatenated row-wise. The (S x A) x d matrix representing a can be given as,
v = M
v
Σ
druf v.
(M* affine for basis s)
V
V
Pslis', 8 ))p((')))))
V Σ mu
= ( (0(())))))0())))
(ov
T $ p
Hence, the successor features are affine with policy independent basis Pyp-
A.5. Proof of Theorem 6.3
Theorem 6.3. If Md (s,ss,, Φ(s, a, st) w and a(s,a,s*) = (o(ss,)TTos(s")) the successor feature 1/" (s, ay"" for the basic feature o. (s)* (0.0) -1
Proof. Consider ((8, a, s+) E Rdd as the set - 1 basis vectors and the bias with w Rd being the d weights to combine the basis and w 1. Clearly from Theorem 4.2, M (s, a, st) can be represented as Φ(s,, a, st w Further, 6(s, st) = @ a) (s+) where (s, igddd and (.(s ) Rd. So.
M" - Sols, () wup
M" ΣΣ
M" = a) oo(
M" = 1d (s
(Writing as) ww as (s,a))
M" = (s,
From Lemma 6.2, v (s, is the successor feature for the basic feature (c())
Note: In continuous settings, we can use the dataset marginal density as described in Section 5. The basic features become
(x)" (E)(
A.6. Deriving basis for the Toy Example
Consider the MDP shown in Figure 5. The state action visitation distribution is written as d
(d(oo, an), dos),ao), d(s0, a1),d(sn) a(s),,))))² The corresponding dynamics can be written as,
#0s
#1, ag
So,
P
80
1
1
0
81
0
1
15

---

Proto Successor Measure:
Representing the Behavior Space of an RL Agent
ag
a1
SO
$1
a1
ag
Figure 5. The Toy MDP described in Section
The Bellman Flow equation thus becomes,
Ed(s, d(s = -))(s)) + + Σ P(s)s',a'ddee d')d(''sa))
d(so: 0o)
d(s0- 2o)
1 o o 0 0 1
d(si, au)
= (1 - 7) (u(Su)
1110 0 1001
d(s1, ao)
d(80, a1)
(p(s))
+7
1
d( a1) d(s1, 01)
d(s1,a1) a1)
d(s0, ap)
1-y
-
d(s1, an)
(1-7) (u(su) (u(s))
7
1 1
d(so;an a1)
d(s1,a1) a1)
This affine equation can be solved in closed form using Gauss Elimination to obtain
d((su ao) d(s1, ao) d(sp, a1) ((s1,a))
1+7 plax 1-yu(
+ W2
1+7
14-
(21)
0
0
1
B. Experimental Details
B.1. Environments
1.1.1. GRIDWORLDS
We use https: /tggtttttbb///g com/ world and 4 room experiments. The task is to reach goal state that is randomly sampled at the beginning of every episode. The reward function 0 at all non-goa non-goal states while at goal states. The episode length for these tasks are 200.
code-base to build upon the grid-
The state representation given by (#, y) which are scaled down be [0, 1]. The action space consists of five actions: {up, right, down, feft, slay).
1.1.2. FETCH
We build on top of https: / / thub. hoom/ wtweeedd-ootti//e contttoocoottolllllllolloo waweerollllllllecmm httpc//onooolllllleeeeeeeeeee which contains the Fetch environments with discretized action spaces. The state space is unchanged but the action space is discretized to produce manhattan style movements i.e. move one-coordinate at a time. These six actions are mapped to the true actions Fetch as: 0..00.,.
{0 : [1, 0, 0, 0j, 1: [o. O, 0). 2: 0, 1, 0) [-1,0,0,0,4 4: of 50.
4: [0, 0, 0], 5 [o, 0, -1, 0]3. Fetch has an episode length
16

---

Proto Successor Measure: Representing the Behavior Space of RL Agent
Figure 6. DM Control Environments: Visual rendering of each of the four DM Control environments we use: (from left to right) Walker, Cheetah, Quadruped, Pointmass
B.1.3. DM-CONTROL ENVIRONMENTS
These continuous control environments have been discussed in length in DeepMind Control Suite (Tassa al., 2018). We these enviromments to provide evaluations for PSM larger and continuous state and action spaces. The following four environments are used:
Walker: It has 24 dimensional state space consisting of joint positions and velocities and 6 dimensional action space where each dimension of action lies in [-1, The system represents a planar walker. At test time, we test the following four tasks: Walk, Run, Stand and Flip, each with complex dense rewards.
Cheetah: has 17 dimensional state space consisting of joint positions and velocities and 6 dimensional action space where each dimension of action lies in [-1, 1]. The system represents planar biped "cheetah". At test time, test the following four tasks: Run, Run Backward, Walk and Walk Backward, each with complex dense rewards.
Quadruped: It has 78 dimensional state space consisting of joint positions and velocities and 12 dimensional action space where each dimension of action lies in [11, The system represents a 3-dimensional ant with legs. At test time, we test the following four tasks: Walk, Run, Stand and Jump, each with complex dense rewards.
Pointmass: The environment represents 4-room planar grid with 4-dimensional state space (xr, M, Uey Fy) and 2- Reach Top Left, Reach Top Right, Reach Bottom Left and Reach Bottom Right each being goal reaching tasks for the four room centers respectively.
dimensional action space. The four tasks that test
All DM Control tasks have an episode length of 1000.
B.2. Datasets
The exploratory data is collected by uniformly spawning the agent and taking random action. Each of the three method is trained on the reward-free exploratory data. At test time, a random goal is sampled and the optimal Q function is inferred by each.
Gridworld:
The exploratory data is collected by running DQN (Mnih al., 2013) training with RND reward (Burda et al., 2019) taken from https: : / github. com/ ccam////rrrggaaaaavveerraaaaaaaa htt.ap/sss...ttttttsssssoooooaaa jectories, each of length 50, are collected.
Fetch:
20000
DM Control: We use publically available datasets from ExoRL Suite (Yarats et al., 2022) collected using RND exploration.
B.3. Implementation Details
B.3.1. BASELINES
We consider variety of baselines that represent different state of the art approaches for zero-shot reinforcement learning. In particular, we consider Laplacian, Berward-Backwarn and HILP.
1. Laplacian (Wu et al., 2018; Koren, 2003): This method constructs a graph Laplacian for the MDP induced by random policy. Eigenfunctions of this graph Laplacian gives a representation for each state Φ(s), or the state feature. These fate-features are used to learn the successor features; and trained to optimize a family of reward functions r(s) = (ø(s) 2). where usually sampled from unit hypersphere uniformly (same for all baselines). The reward functions are optimized
17

---

Proto Successor Measure: Representing the Behavior Space of RL Agent
via TD3.
2.
(Blier et al., 2021a; Touati & Ollivier, 2021a; Touati et al., 2023); gerward-backweer algorithm takes slightly different perspective: instead of training state-representation first, a mapping is defined between reward function to a latent variable (z = E. (s)).r(s) and the optimal policy for the reward function is set to T.. i.c the policy conditioned on the corresponding latent variable 2. Training for optimizing all reward functions in this class allows for state-features and successor-featuree to coemerge. The reward functions are optimized via TD3.
HILP (Park et al., 2024a): Instead letting the state-features coemerge as FB, HILP proposes to learn features from offline datasets that sufficient for goal reaching. Thus, two states are close each other if they are reachable in few steps according to environmental dynamics. HILP uses specialized offline RL algorithm with different discounting to learn these state features which could explain its benefit in some datasets where TD3 is not suitable for offline learning.
Implementation: We build upon the codebase for FB https:/ / i thub. com/ facebookresearchh controllable_agee and implement all the algorithms under a uniform setup for network architectures and same hyperparameters for shared modules across the algorithms. We keep the same method agnostic hyperparameters and use the author-suggested sethod-specffc hyperparameters. The hyperparameters for all methods can be found here:
Table Hyperparameters baselines PSM.
Hyperparameter
Value
Replay buffer size
5 x 10t (10 10 for maze)
Representation dimension
128
Batch size
1024
Discount factor
0.98 (0.99 for maze)
Optimizer
Adam
Learning rate
x 10
Momentum coefficient for target networks
0.99
Stddev for policy smoothing
0.2
Truncation level for policy smoothing
0.3
Number of gradient steps
2 x 10
Batch size for task inference Regularization weight for orthonormality loss (ensures diversity)
104
FB specific hyperparameters
Hidden units (F)
1024
Number of layers (P))
Hidden units (b)
256
Number of layers (b)
2
HILP specific hyperparameters
Hidden units (Φ)
256
Number layers (Φ)
Hidden units (1b)
1024
Number of layers (v))
Discount Factor for
0.96
Discount Factor for //
0.98 (0.99 for maze)
Loss type
Q-loss
PSM specific hyperparameters
Hidden units (ø. b)
1024
Number of layers (Ø, b)
Hidden units (w)
1024
Number of layers (w)
Double GD Ir
lc-4
Proto Successor Measures (PSM): PSM differs from baselines in that we lear richer representations compared to
Laplacian or HILP as we are not biased by behavior policy or only learn representations sufficient for goal reaching. Compared FB, our representation learning phase more stable as we learn representations by Bellman evaluation backups
18

---

Proto Successor Measure: Representing the Behavior Space of an RL Agent
and do not need Bellman optimality backups. Thus, our approach is susceptible learning instabilities that arise from overestimation that is common in Deep RL and makes stabilizing FB hard. hyperparameters are discussed in Appendix Table 2.
19

---

Proto Successor Measure: Representing the Behavior Space of RL Agent
B.3.2. PSM REPRESENTATION LEARNING PSUEDOCODE
def psm_loss( self, obs: terch.Iensox action: torch,Tenson
discount: torsh. next_obs: torch. next_goal; torch. z: Terch.Tensor,
step: int
Ap.Dict[str, (str, float]:
metrics: sp.Dict[sti. float]
12
Create a batch_size batch_size for learning M (")pi
13
idx
shape (0]) to (obs device)
14
mesh = torch.stach (torch.meshgrid
idx, indexing= (xy re.rase (-1, 2)
m_cbs cbs [mesh 0]]
16
n_next_cbs next_obs [mesh [m, 0]]
17
mactionn action (meshli,, 0]]
18
n_next_godd goal = next goal (meshl;, 1 perm = torch. randperm (obs shape [0])
19
compute loss with terch.no_grade target_phal target_b sewf. target (m _e_eet__bba m __nnet_gooll
22
23
target_w self. self.w_target (z)
target_phi target_phi arange (target_ (targggg shape next_actions.
squeeze (1)]
26
target_b =
next_actions. equuees (1))
27
target_M. = forch.cinsume
sd -> s", target_phi, target_w) + target_b
10
phi, b self.psm (m_obs, m_next_goal)
31
phi philtorch. spai.shape m_action.squest (1)
32
= iterchharange (b.
torch.einsum""sd. shapel01),
squeeze (1)
s", phi, self.x(z))
34
M. (
(obs
shspshap(001)
35
target_M target_M. reshape (obs. shapee [0], obs obs.shape[0]]
36
(MM. off _tf__daag = "I.bool()
- torch.eye (MM. size (), Sevice-K.devica.
psm_ offiiigd:::: tp.Any 0.5 (M discount target_M, M) [off diag] ppo mean .nen ()
39
psm_diags diag: An. ((1 discount) * (M.diag () unsqueezze .nea)
40
pem_loss = psm_offdias psm_diag
41
optimize PSM
44
sel st_ self self.actor_opttzerr..o pam loss.backwarr)))
45
(set_to_none-
46
self. self.opt.step self. opt step
48
All experiments were trained Intel(R) Xeon(R) CPU E5-2620 v3 @ 2.40GHz CPUS and NVIDIA GeForce GTX TITAN GPUs. Each training took around 10-12 hours.
Compute:
C. Additional Experiments
C.1. Ablation on dimension of the affine space: d
We perform the experiments described in Section 7.3 for two of the conitnuous environments with varying dimensionality of the affine space (or corresponding successor feature in the inductive construction), d. Interestingly, the performance of PSM does not change much across different values of d ranging from 32 to 256. This is in contrast to methods like HILP which sees significant drop in performance by modifying d.
20

---

Proto Successor Measure: Representing the Behavior Space of RL Agent
Environment
Task
32
50
128
Walker
Stand
898, 98 + 48.64
942.85 + 19.43
872.61
+ 38.81
911.25 +: 32.86
Run
359.51 70.66
392.76 + 31.29
351.50
19:46
372.39 41.29
Walk
825.66 60.14
822.39 1: 60.14
891 44 46.81
886.03 28.96
Flip
628.92 + 94,95
521.78 + 29.06
640.75 = 31.88
593.78 + 27.14
Average
678.27
669.45
689.07
690.86
Cheetah
Run
298.95 988.98 1. 95.63
386.75 1. 55.79
276.41 1: 70.23
268.91 1: 79.07
Run Backward
295.43 + 19.72
260. 1: 24.93
286. +: 25. 38
290.89 1: 14.36
Walk
12 84.25
893.89 89 + 91.69
887.02 + 59.87
920.50 98
Walk Backward
978.64 # 8.74
916.68 124.34
980,90 + 2.04
982.29 = 0.70
Average
628.79
615.61
607.61
615.64
Table 3. Table shows comparison (averaged over 5 seeds) hetween different representation sizes (or affine space dimensionality d) for PSM.
C.2. Quantitative Results on Gridworld and Discrete Maze
We provide quantitative results for experiments performed in Section 7.1.
Quantitative Experiment Description: For each randomly sampled goal, obtain the inferred value function and the inferred pol- icy using PSM and the baselines. At every state in the discrete space, we check if the policy inferred by these algorithms is optimal or not. The oracle or the optimal policy can be obtained by running the Bellman Ford al- gorithm in the discrete gridworld maze. We report (in Table 4) the average error incorrect policy predictions/Total # of states) for 10 randomly sampled goal (over 3 seeds).
Enviromment
Laplace
H-B-biased
PSM
Gridworld
19.28 2.34
14.53 4: 0.68
2.13 4: 0.73
2.05 1: 1.20
38.47
7.01
28.80
12.9 1.63
11.54 1.07
Table Table shows average error (averaged seeds) for the predicted policy from different zero-shot RI methods with respect to the oracle optimal policy.
The original FB method assumes that tasks are goal-conditioned, implicitly biasing the training to smaller set of reward functions. We remove that bias in our gridworld experiments to ensure an apples-to-apples comparison. We present the results for the biased goal-condi sampling as "FB-biased" in Table 4.
As clearly seen, the average error for PSM is significantly less than the baselines which augments the qualitative results presented in Section 7.1.
21


Reviews:
[{"review_id": "c7S6cDmG5T", "reviewer": "Reviewer_MScc", "summary": "This paper studies the reward-free RL problem and proposes the concept of proto successor measure (PSM), which is built on the idea of proto value functions in (Mahadevan and Maggioni, 2007) and serves as the basis set of all the possible visitation distributions in a given RL environment. By learning this basis set, given a reward function at the inference stage, solving for an optimal policy is equivalent to solving a constrained linear program, where both the objective function and the Bellman flow constraint can be expressed in terms of this basis. Accordingly, this paper proposes to implement the PSM by integrating two components: (1) Using the measure loss similar to that in (Touati & Ollivier, 2021), one can learn the weight that corresponds to a given policy under some learned PSM. (2) To learn the basis set in PSM via pre-training, this paper proposes to construct a discrete codebook of policies, which serves as a way to simulate uniform sampling over policies. \n\nThe proposed PSM is evaluated on both discrete and continuous control environments, including gridworld-like problems, Fetch-Reach tasks, and the ExoRL suite. Through experiments, the PSM is shown to be comparable or outperform the recent benchmark reward-free RL methods.\n\n## update after rebuttal\n\nI appreciate the authors for the additional response and for answering my follow-up questions. The additional sanity check of FB and the clarification on linear mapping are helpful.\n\nOn the other hand, I am still not convinced by the explanation about overestimation bias from the off-policy learning perspective (FB in (Touati et al., 2023) was implemented based on TD3, and there are existing techniques to largely avoid overestimation) and the argument “...This instability is absent in PSM as PSM performs off-policy evaluation using samples from the codebook policy…”\n\nAfter going through all the reviews and the rebuttal responses again, overall, the idea of PSM appears to be a useful addition to the reward-free RL literature with good empirical support. Hence, I have increased my score accordingly and lean towards acceptance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback. We would like to address all the concerns of the reviewer below:\n\n> It would be good to compare a more diverse set of reward functions for each environment.: \n\nWe 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.\n\n\n> Comparison with successor feature methods: \n\nWe 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. \n\n> However, this is quite different from the results provided in the original FB paper. It is unclear where the differences come from. : \n\nWe 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. \n\n\nAdditionally, 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.\n\n> 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. : \n\nWe 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:\n\n\n(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.\n\n\n(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.\n\n\nIn contrast, PSM avoids both assumptions and directly aiming to accurately capture all feasible visitation distributions. \n\n> It remains unclear to me why PSM is a better approach than the FB representation: \n\n(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.  \n\n> 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. : \n\nWe 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.\n\n> An ablation study on the need for the discrete policy codebook would be helpful.: \n\nThe 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.", "claims_and_evidence": "The main theoretical claims of this paper (Theorems 4.1 & 4.4, Corollaries 4.2-4.3, and Theorems 6.1 & 6.3) are supported by the proofs in the appendix.", "methods_and_evaluation_criteria": "- Like the existing reward-free RL literature, this paper evaluates the capability of zero-shot generalization by testing the performance of PSM and other baselines across a few difference reward functions (e.g., for the Walker environment, take the reward functions of the four underlying tasks Stand, Run, Walk, and Flip). One concern is that it is unclear whether this small set of reward functions is already representative enough. While I know this evaluation procedure is also adopted by some prior works (e.g., (Touati & Ollivier, 2021) and (Touati et al., 2023)), to make a more thorough comparison, it would be good to compare a more diverse set of reward functions for each environment. \n\n- All the pre-training is done on a shared offline dataset constructed using a fixed behavior policy, and this looks fair to all the algorithms.", "theoretical_claims": "- The main theoretical results of the paper are two-fold: (1) The set of possible visitation distributions form an affine set (Section 4); (2) The classic successor features can also be represented using a similar basis set. \n\n- I have checked the proofs in the appendix up to Appendix A.3, and I did not spot any major issues regarding the correctness of the analysis.", "experimental_designs_or_analyses": "I have checked the experimental designs in Section 7. Most of them look reasonable. \n\nHowever, there are two things that require further explanation:\n- Representative reward functions at test time: As mentioned above, one concern is on whether the small set of reward functions considered in Section 7 for each environment is indeed representative enough. Specifically, like the existing reward-free RL literature, this paper evaluates the capability of zero-shot generalization by testing the performance of PSM and other baselines across a few difference reward functions (e.g., for the Walker environment, take the reward functions of the four underlying tasks Stand, Run, Walk, and Flip). While I know this evaluation procedure is also adopted by some prior works (e.g., (Touati & Ollivier, 2021) and (Touati et al., 2023)), to make a more thorough comparison, it would be good to compare a more diverse set of reward functions for each environment. \n\n- Comparison with successor feature methods: The PSM approach is highly related to the general method of Successor Features (SF). Therefore, it would be helpful to compare PSM with SF-based approaches, e.g., (Alegre et al., 2022) to demonstrate the benefits of learning this basis set. \n\n- Regarding the qualitative results in Figure 3, it appears that FB cannot generalize well in the four-room environment. However, this is quite different from the results provided in the original FB paper (Touati & Ollivier, 2021). It is unclear where the differences come from. It would be good to describe this in more detail to ensure a fair comparison.\n\n\n\n(Alegre et al., 2022) Lucas N. Alegre, Ana L. C. Bazzan, and Bruno C. da Silva, “Optimistic Linear Support and Successor Features as a Basis for Optimal Policy Transfer,” ICML 2022.", "relation_to_broader_scientific_literature": "The main contribution of this paper is to offer another perspective of solving reward-free RL by pre-training a basis set called Proto Successor Measure. While this line of research can be traced back to the classic proto value functions (Mahadevan and Maggioni, 2007), this paper offers a practical implementation that scales the idea of proto value function to a broader class of RL problems, including high-dimensional robot control. Given the recent attention on RL foundation models, I find this paper to be quite nice in contributing to the general literature of RL foundation models.", "essential_references_not_discussed": "As far as I know, most of the recent relevant works on learning reward-free RL are cited in this paper. That being said, to better describe the contributions of this papers, there are two things to clarify:\n\n- Comparison with successor feature methods: 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. Theoretically speaking, using SF, like in (Lehnert & Littman, 2020; Hoang et al., 2021; Alegre et al., 2022; Reinke & Alameda-Pineda, 2021), can already be sufficient to achieve similar zero-shot generalization.\n- Comparison with FB representation: While I can appreciate the idea of PSM, it remains unclear to me why PSM is a better approach than the FB representation. Indeed, as also shown in Table 1, the two approaches are just comparable in most of the tasks. Moreover, compared to FB, PSM requires solving an additional constrained linear program to get a policy at inference time. Accordingly, it appears that FB could be preferred in practice. While it is mentioned in Lines 155-163 that “FB ties the reward with the representation of the optimal policy derived using Q function maximization which can lead to overestimation issues and instability during training as a result of Bellman optimality backups,” I do not find this argument convincing. It would be helpful to provide more evidence to justify this argument.", "other_strengths_and_weaknesses": "All the strengths and weaknesses are provided in the sections above.", "other_comments_or_suggestions": "- 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.\n- Equation (3) can be reorganized. There are two separate “such that” which can be merged into one.\n- An ablation study on the need for the discrete policy codebook would be helpful.", "questions_for_authors": "Please see the above questions on the comparison with SF-based and FB-based methods, the baselines for evaluation, as well as the experimental configuration for the test-time evaluation. I will be willing to reconsider my evaluation if the authors can address these issues.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LKfFl4FsDq", "reviewer": "Reviewer_Gftv", "summary": "The paper investigates representation learning in RL with the aim of performing zero-shot learning: computing optimal policies on downstream tasks wihout any further training. Building on earlier works about representation learning in a reward-free setting, especially on successor representations, it proposes a new representation called \"proto-successor measure\" (PSM) and gives a procedure to learn this representation. It is based on the observation that successor measures (the discounted distribution of states under a given policy) satisfy a Bellman equation which is independent of the policy. As solutions of linear systems, successor measures and related objects thus span an affine subspace. The proto-successor measure consists then in computing an approximation of this affine subspace. This theory is completed with experimental works in gridworld and continuous control environments, exhibiting better performance than forward-backward, HILP and Laplacian representations, which are popular representations that attracted a lot of attention recently.\n\n## Update after rebutal:\nI think the authors for addressing my last questions. I agree with the authors that PSM may bring a complementary and novel idea to the FB representation of Touati \\& al, so I will maintain my score (weak accept).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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:\n\n> I think the \"mathematically complete\" mention is a bit exaggerated here. The paper would require a much more quantitative analysis to deserve that mention.: \n\nBy “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.\n\n> I am not sure think PSM really compares with FB, Laplacian and the likes: \n\nWe 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.\n\n> Typos: \n\nWe thank the reviewer for pointing out the typos. We shall correct them in the camera-ready version.\n\n> I think the theoretical content is quite simple so I would suggest putting the focus on the experimental work.: \n\nWe 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. \n\n\n> 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.: \n\nWe 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.\n\n\n> 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?: \n\nAs 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.", "claims_and_evidence": "The authors claim their work on the PSM representation gives a \"novel, mathematically complete perspective\" on representation learning in RL, together with an efficient algorithm. While I think the present work has an interest, I am a bit more reserved about its ground-breaking aspect.\n\nFrom a theoretical perspective, I think the main contribution is in the observation that there exists a Bellman equation independent of the policy. This is a simple fact which could nonetheless have powerful consequences and to the best of my knowledge this is indeed new. However I don't think the rest of the theoretical content goes much deeper than this: the main result Theorem 4.1 and its corollaries for instance are simple statements that solutions of a linear system span an affine subpsace. So I think the \"mathematically complete\" mention is a bit exaggerated here. The paper would require a much more quantitative analysis to deserve that mention.\n\n\nWhen it comes to the experiments, the authors underline the apparently better performances of PSM over other representations like the forward-backward (FB) model of Touati \\& al, but I am not sure this comparison is really fair. I believe the generalization power of PSM is a bit different: as the authors write in Theorem 6.3, the PSM representation factorizes the successor measure as $M^{\\pi}(s,a,s^{+}) = \\phi(s,a,s^{+}) w^{\\pi}$, while the FB model adds a further factorization in that $M^{\\pi}(s,a,s^{+}) = \\psi^{\\pi}(s,a) \\phi(s^{+})$ (I consider a tabular setting). Therefore I think PSM allows a generalization over policies but it remains intrinsically a high-dimensional ($S^2 A$), more complex representation than FB which generalizes both over policies and state-actions. Thus it hardly comes to as a surprise that PSM provides more expressive representations than all other benchmarks considered here, and I suspect that in a tabular a setting the sample-complexity of learning PSM could scale as $\\Theta(S^2 A)$.", "methods_and_evaluation_criteria": "The overall method to learn the PSM representation makes sense. I believe however there is a typo in the loss function p.5, I think there is a square as well as a $(1-\\gamma)$ factor missing. Also I haven't quite understood the idea of the discrete codebook of policies, and in particular why the \"approach provably samples from among all possible deterministic policies\".\n\nFor evaluation criteria I don't have concerns other than the one above: I am not sure think PSM really compares with FB, Laplacian and the likes.", "theoretical_claims": "I haven't checked all the proofs but I think the results are simple enough to be correct.", "experimental_designs_or_analyses": "I don't have more to say than in \"Methods and Evaluation Criteria\".", "relation_to_broader_scientific_literature": "The paper seems to be heavily inspired by works on Laplacian representations as well as the recent line of work of Touati \\& al who proposed the forward-backward model representation to perform zero-shot learning. The present paper starts by considering the same object, the successor measure, which for a fixed policy is the operator that maps a reward to its value function. It has appeared quite often in the literature. The forward-backward model of Touati \\& al. aims to compactly represent the successor measure of all optimal policies though what is essentially a low-rank matrix in the tabular setting. The PSM representation bears resemblance in that it attempts to represent successor measures for different policies, but the representation is now low-rank. It leverages the fact that successor measures span an affine subpsace, which I think has not been observed in the work of Touati \\& al.", "essential_references_not_discussed": "I think the paper already discusses the most important references in relation to the subject.", "other_strengths_and_weaknesses": "The main idea of the paper is simple but interesting. The paper is in general pretty clear.", "other_comments_or_suggestions": "Overall I think the basic idea of this paper, that successor measures span an affine subspace, has a good potential and the experiments are promising. However I think the theoretical content is quite simple so I would suggest putting the focus on the experimental work. For instance, I don't think Section 3.2 on affine spaces rally brings a lot to the paper. I would assume most researchers in machine learning know what an affine space is.\n\nSecondly, 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.\n\nFinally I have noticed a few typos here and there: \"familia r' line 043, \"ideplogy\" p.8", "questions_for_authors": "1. I'd like primarily the authors to address the concern raised above: 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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fZZWAJahSZ", "reviewer": "Reviewer_UwHe", "summary": "The paper introduces Proto Successor Measure (PSM), a basis set to represent all possible behaviors of an RL agent in an environment. The key insight is that any valid successor measure must satisfy the Bellman flow equation. By rearranging the Bellman flow equation one gets an affine equation. Any solution to this affine equation can be represented as an affine combination of a basis set, where the bases and the bias term are independent of the policy. Hence, by learning this basis set, one can represent the successor measure of any policy by finding the corresponding set of linear weights $w$. Given a downstream reward, one can solve a linear program to obtain the $w$ corresponding to the successor measure for the optimal behavior, obtain the optimal Q function, and then back out the optimal policy. With these theoretical foundations, the paper introduces a practical algorithm to learn the basis functions using reward-free interaction data. This involves using a seeded RNG to draw samples from the policy space. The method is evaluated on a selection of manipulation and locomotion zero-shot RL problems from FetchReach and ExoRL, where PSM exhibits superior transfer behavior compared to prior zero-shot RL baselines.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer for providing detailed feedback on our paper. We would like to clarify all the questions raised by the reviewer:\n\n> The inference for the linear weights involves solving a linear program with constraints. This is harder to solve than linear regression in prior works. : \n\nYes, 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. \n\n\n> 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. : \n\nFollowing 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.\n\n\n> 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:\n\nThe 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. \n\n> Typos: \n\nWe thank the reviewer for pointing out these typos. We shall fix them in the camera-ready version. \n\n> How do you solve (4) in continuous space? Do you sample (s, a) from a dataset? \n\nIn both continuous and discrete spaces, we sample (s, a) from the dataset. \n\n> How do you determine the dimensionality of the basis set? Can you provide an ablation of the dimensionality of the bases? \n\nWe 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].\n\n[1]: Touati, Ahmed, and Yann Ollivier. \"Learning one representation to optimize all rewards.\" Advances in Neural Information Processing Systems 34 (2021): 13-23.\n\n> In hyperparameters (Table 2), why does w have 3 layers? Is it represented by a neural network? \n\nDuring 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$.", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence. Empirically, the paper claims that PSM leads to more desirable value functions and better transfer performance. The visualizations of value functions in the Gridworld and Four-room environments (Figure 3) show that PSM value functions are more concentrated than the FB and Laplace value functions. On FetchReach and ExoRL (Figure 4, Table 1), PSM achieves higher episode returns than baselines when transferring to downstream rewards. For theoretical claims, please refer to the discussion in the \"Theoretical Claims\" section.", "methods_and_evaluation_criteria": "This paper applies PSM to the problem of zero-shot reinforcement learning, which aims to learn a representation of policies from a reward-free dataset and quickly obtain the optimal policy for some test-time reward functions without additional environment interaction. The evaluation domains properly represent this problem, where the methods are trained on some large reward-free datasets and transferred to multiple different reward functions.", "theoretical_claims": "I checked the correctness of the proofs for the theoretical claims. The paper provide three sets of theoretical claims:\n1.  (Theorem 4.1, Corollary 4.2, Corollary 4.3) Deriving proto successor measure as the basis set for solutions to an affine equation.\n2. (Theorem 4.2) The basis value functions represent a smaller space of value functions than proto successor measures. \n3. (Theorem 6.1, Lemma 6.2, Theorem 6.3) Connects PSM to successor features showing that one can decompose PSM into successor features corresponding to some state features. These proofs largely follow the mechanisms in [1].\n\nThe proofs for these claims appear to be sound.\n\n[1] Ahmed Touati, Jérémy Rapin, Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? ICLR 2023.", "experimental_designs_or_analyses": "I checked the soundness of the experimental designs, including the Gridworld analysis, the FetchReach environment, and the ExoRL benchmark.", "relation_to_broader_scientific_literature": "The key contributions are built on a line of prior works on zero-shot reinforcement learning. Particularly relevant is the FB representation [1,2], which proposes to decompose a successor measure into a product of a forward representation and a backward representation. Unlike FB where the representations are conditioned on some policy representation $z$, PSM learns policy-independent basis functions, which is more desirable.\n\n[1] Ahmed Touati, Yann Ollivier. Learning One Representation to Optimize All Rewards. NeurIPS 2021.\n[2] Ahmed Touati, Jérémy Rapin, Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? ICLR 2023.", "essential_references_not_discussed": "The paper includes sufficient context for the reader to understand its contributions.", "other_strengths_and_weaknesses": "**Strengths**\n1. This work proposes a principled approach to zero-shot RL, a problem of great interest to the community.\n2. The derivations are driven by first principles and mathematically sound. The practical implementation achieves superior empirical performance compared to relevant baselines.\n3. This work improves upon prior zero-shot RL methods by learning representations truly independent of the policy and in principle capable of representing all behaviors in an RL environment.\n\n**Weaknesses**\n1. The inference for the linear weights involves solving a linear program with constraints. This is harder to solve than linear regression in prior works.\n2. 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. \n3. 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.", "other_comments_or_suggestions": "1. Line 43 typo \"familia r\"\n2. Line 89 typo \"ge neralization\"\n3. Line 177 \"Successor measures are more general than state-action visitation distributions as they encode the visitation of the policy conditioned on a starting state-action pair.\" This statement is inaccurate because successor measures are a special case of state-action visitation distributions (equation 2). \n4. Line 189 \"we do not need the first constraint on the linear program (in Equation 3) anymore\" -> the second constraint.", "questions_for_authors": "1. How do you solve (4) in continuous space? Do you sample (s, a) from a dataset?\n2. How do you determine the dimensionality of the basis set? Can you provide an ablation of the dimensionality of the bases?\n3. In hyperparameters (Table 2), why does w have 3 layers? Is it represented by a neural network?\n\n## Updates After Rebuttal\nThe authors have adequately addressed my questions. I will maintain my score.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "g29MnRfaxe", "reviewer": "Reviewer_PPqR", "summary": "The paper constructs a linear framework, \"Proto Successor Measures\", for classifying the space of Q functions, a generalization of successor features. The paper provides some theoretical results on PSM, showing how they can be learned from offline data and used for inference at test-time. Experiments show the advantage of PSM over Laplacian and Forward-Backward baselines on gridworlds and locomotion tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for providing detailed feedback on our paper. We would like to address the concerns raised by the reviewer:\n\n> At best, further samples are needed from the environment to infer a reward function. This could be problematic especially in sparse environments. : \n\nWe 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.\n\n>Theorem 4.4: The notation is somewhat confusing for me here. Fixing typography  can also help readability.: \n\nWe 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. \n\n> The proof seems to follow with a straightforward rearrangement of terms, but it is a bit opaque (to me): \n\nThe 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.\n\n\n> Fig 4: What are update steps?  seeds is quite small -- adding more could strengthen the statistics.: \n\nUpdate 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. \n\nWe appreciate the reviewer for pointing this out and will add more seeds during the camera ready version.\n\n> Table 4: Do you know how such errors scale with size of maze?: \n\nThere 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. \n\n> How is there any \"smoothness\" between goals in the DMC tasks? They seem quite distinct, so how is there a generalization effect among tasks? \n\nWe 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.\n\n> Could you further clarify the use/meaning of the bias vector and how its presence departs from typical SF?: \n\nThe 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.\n\n> 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?:\n\nThe 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. \n\n\n> Can you provide an experiment showing the result of small datasets/dataset scaling? : \n\nWe 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.", "claims_and_evidence": "Claim: Visitation distributions lie on a convex hull defined by affine constraints.\n\nEvidence: Theorem 4.1, 4.2.\n\nClaim: Better generalization properties\n\nEvidence: Theorem 4.4, Fig 1 table 1.\n\nClaim: The optimal policy can be produced for any reward function.\n\nEvidence: This seems over-stated. At best, further samples are needed from the environment to infer a reward function. This could be problematic especially in sparse environments.", "methods_and_evaluation_criteria": "The discrete codebook seems practically relevant and improves efficiency. The evaluation on diverse environments showcases the efficacy of the method in a variety of settings.", "theoretical_claims": "Many results stem from the linearity of the framework. Making use of the Bellman flow equation is nice.\n\nTheorem 4.1: Its proof based on the Bellman flow equation and related arguments appears sound.\n\nCor 4.2: The affine set proof follows straightforwardly.\n\nTheorem 4.4: The notation is somewhat confusing for me here. E.g. maybe there are missing brackets around $span \\\\{ \\Phi^{vf} \\\\}$? Fixing typography $span \\to \\textrm{span}$ can also help readability. The proof seems to follow with a straightforward rearrangement of terms, but it is a bit opaque (to me): can you provide some further intuition on the terms like $\\beta^\\pi , k$, and their ratio? L692 is a bit confusingly worded I feel. Do you mean to say that you've shown the same $V^\\pi$ is represented in $\\\\{ span \\\\{\\Phi \\\\} r \\\\}$? The construction of $\\Phi(s,s')$ from the indicator function seems trivial, but I might be missing something deeper here. Some further explanation would be appreciated! If I understand correctly, the proof is highlighting that only factorized $\\Phi$ lead to value functions represented in the smaller span?\n\n\n- Lets $\\to$ Let's\n- L709, Value $\\to$ value", "experimental_designs_or_analyses": "Experiments are aligned with the claims.  \n- Visualization of the Q function in Fig 3 is nice, but how close is this to the optimal Q function?\n- Fig 4: What are update steps? $3$ seeds is quite small -- adding more could strengthen the statistics.\n- Table 3 helps understand different representation sizes (based on comments in final section)\n- Table 4: Do you know how such errors scale with size of maze?\n- How is there any \"smoothness\" between goals in the DMC tasks? They seem quite distinct, so how is there a generalization effect among tasks?", "relation_to_broader_scientific_literature": "This paper nicely positions PSM as a \"successor\" to the successor features framework, which has proven itself in recent years. Generalizing the basis with policy-independence is a useful step forward for such zero-shot solutions and understanding the MDP geometry.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "N/A (please see other comments)", "other_comments_or_suggestions": "- Could you further clarify the use/meaning of the bias vector and how its presence departs from typical SF?\n- 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?\n- Can you provide an experiment showing the result of small datasets/dataset scaling?", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In this work the authors study the problem of representation learning in reinforcement learning tasks where they develop methods for zero shot learning. Based on the observation that successor measures satisfy a Bellman equation that is independent of the policy they introduce the Proto Successor Measure (PSM). This construction is based on the observation that any solution to an affine equation can be represented as an affine combination of a basis set, and therefore one can represent a successor measure by a corresponding set of linear weights. This work’s conceptual and theoretical contributions are accompanied with experiments in gridworld and control environments showing the benefits of their approach.
Decision: Accept (poster)