Distributional Successor Features Enable Zero-Shot Policy Optimization
Authors: Chuning Zhu
Xinqi Wang
Tyler Han
Simon Shaolei Du
Abhishek Gupta
Venue: NeurIPS 2024 poster (2024)
Keywords: Deep Reinforcement Learning
Multitask Transfer
Diffusion Models
Primary area: reinforcement_learning

Abstract: Intelligent agents must be generalists, capable of quickly adapting to various tasks. In reinforcement learning (RL), model-based RL learns a dynamics model of the world, in principle enabling transfer to arbitrary reward functions through planning. However, autoregressive model rollouts suffer from compounding error, making model-based RL ineffective for long-horizon problems. Successor features offer an alternative by modeling a policy's long-term state occupancy, reducing policy evaluation under new rewards to linear regression. Yet, policy optimization with successor features can be challenging. This work proposes a novel class of models, i.e., Distributional Successor Features for Zero-Shot Policy Optimization (DiSPOs), that learn a distribution of successor features of a stationary dataset's behavior policy, along with a policy that acts to realize different successor features within the dataset. By directly modeling long-term outcomes in the dataset, DiSPOs avoid compounding error while enabling a simple scheme for zero-shot policy optimization across reward functions. We present a practical instantiation of DiSPOs using diffusion models and show their efficacy as a new class of transferable models, both theoretically and empirically across various simulated robotics problems. Videos and code are available at https://weirdlabuw.github.io/dispo/.

Distributional Successor Features Enable Zero-Shot Policy Optimization
Chuning Zhu
Xinqi Wang University of Washington wxqkaxddecs washington. edu
University of Washington zchuninggcs. washington edu
Tyler Han
University of Washington
than123@cs washington, edu
Simon Shaolei Du University of Washington ssdu@cs. washington. edu edu
Abhishek Gupta University of Washington abhguptaccs. washington edu
Abstract
Intelligent agents must be generalists, capable of quickly adapting to various tasks. In reinforcement learning (RL), model-based RL learns dynamics model of the world, in principle enabling transfer to arbitrary reward functions through planning. However, autoregressive model rollouts suffer from compounding error, making model-based RL ineffective for long-horizon problems. Successor features offer alternative by modeling policy's long-term state occupancy. reducing policy evaluation under new rewards to linear regression. Yet, policy optimization with successor features can challenging. This work proposes a novel class of models, i.e., Distributional Successor Features for Zero-Shot Policy Optimization (DiSPOs), successor features of stationary dataset's behavior policy, along with a policy that acts to realize different successor features within the dataset. By directly modeling long-term outcomes in the dataset, DiSPOs avoid compounding error while enabling simple scheme for zero-shot policy optimiza- tion across reward functions. We present practical instantiation of DiSPOs using diffusion models and show their efficacy theoretically and empirically across various simulated robotics problems. Vidcos and code available httttss://wiiddeeeee gi github. 1o/ 1al.oooo/ssss/ dispo/
that learn a distribution
class of transferable models, both
Introduction
Reinforcement learning (RL) agents are ubiquitous in a wide array of applications, from language modeling [8] to robotics [22 28) Traditionally, RL has focused the single-task setting, learning behaviors that maximize a specific reward function. However, for practical deployment, RL agents must be able generalize across different reward functions within environment. For example, robot deployed in a household setting should not be confined to single task such as object relocation but should handle various tasks, objects, initial and target locations, and path preferences.
This work addresses the challenge of developing RL agents that can broadly generalize to any task in environment specified by reward function. To achieve this type of generalization, consider the paradigm of pretraining on an offline dataset of transitions and inferring optimal policies for downstream tasks from observing task-specific rewards. Since the target task not revealed during
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Offline Training of Distribution Successor Features
Adaptation to Downstream Tasks
What are the possible long-term outcomes in this environment?
Cittaee: Uacabaled Datane
Downstream Task
DiSPOs
=(als,
How do I act to accomplish particular outcome W/?
Downstream Task
Figure 1: The transfer setting for DiSPOs. Given an unlabeled offline datasct, DiSPOs model both "what happen?" p(U)s) and "how can we achieve particular outcome?" p(a vb). This used for quick adaptation to new downstream tasks without test-time policy optimization.
pretraining, the model must encode information about the environment dynamics without committing to particular policy reward. Moreover, once the task reward is observed, the model must provide way to quickly evaluate and improve the policy since different tasks require different optimal policies.
A natural approach to this problem is moodel-based reinforcement learning [6] 21 62
which learns an approximate dynamics model of the environment. Given downstream reward function, task- optimal behavior can be obtained by "planning" via model rollouts (50 58
39
45
Typically, model
rollouts are generated autoregressively, conditioning each step on generation from previous step. In practice, however, autoregressive generation suffers from compounding error [31] [ 26], which arises when small, one-step approximation errors accumulate over long horizons. This leads to rollout trajectories that diverge from real trajectories, limiting many model-based RL methods to thort-horizon low-dimensional problems.
An alternative class of algorithms based on successor features (SFs) has emerged as a potential approach to transferable decision-main making 2]. Successor features represent the discounted sum of features for given policy. Assuming linear correspondence between features and rewards, policy evaluation under new rewards reduces to simple linear regression problem. Notably, by directly predicting long-term outcomes, SFs avoid autoregressive rollouts and hence compounding error. However, the notion successor features deeply tied the choice particular policy. This policy dependence hinders the recovery of optimal policies for various downstream tasks. Current approaches to circumvent policy dependence either maintain set of policies and select the best one during inference (3] randomly sample reward vectors and make conditional policy improvements 5,54 55 Nevertheless, turnkey solution to transfer remains a desirable goal.
In this work, we propose new class of models- D Successor Features for Zero-Shot Policy Optimization (DiSPOs) that are rapidly transferable across reward functions while avoiding compounding errors. Rather than modeling the successor features under particular policy, DiSPOs model the distribution of successor features under the behavior policy, effectively encoding all possible outcomes that appear in the dataset at each state. Crucially, by representing outcomes as successor features, we enjoy the benefit of zero-shot outcome evaluation after solving a linear reward regression problem, without the pitfall of compounding error. In addition to the outcome model, DiSPOs jointly learn a readout policy that generates an action to accomplish a particular outcome. Together, these models enable zero-shot policy optimization [55] for arbitrary rewards without further training: at test time, we simply perform a linear regression, select the best in-distribution outcome, and query the readout policy for action to realize it. DiSPOs are new class of world models as they essentially capture the dynamics of the world and can be used to plan optimal actions under arbitrary rewards, without facing the pitsalls of compounding error.
Since multiple outcomes can follow from a particular state, and multiple actions can be taken to achieve particular outcome, both the outcome distribution and the policy require expressive model

---

classes to represent. We provide practical instantiation of DiSPOs using diffusion models 223 47 and show that under this parameterization, policy optimization can be cast as a variant of guided diffusion sampling 13 We validate the transferability of DiSPOs across suite of long-horizion simulated robotics domains and further show that DiSPOs provably converge to "hest-in-data" policies. With DiSPOs, we hope to introduce a new way for the research community to envision transfer in reinforcement learning, and to think of alternative ways to address the challenges of world modeling.
2 Related Work
Our work has connections to numerous prior work on model-based RL and successor features.
Model-Based RL To enable transfer across rewards, model-based RL learns one-step (or multi-
step) dynamics models via supervised learning and use them for planning [I 38 39 58 or
policy optimization 12 50 26
64
20]. These methods typically suffer from compounding error,
where autoregressive model rollouts lead to large prediction errors over time [T 31.. Despite
improvements to model architectures [20] 22, [ 32 65 and learning objectives (27), modeling
over long horizons without compounding error remains an open problem. DiSPOs instead directly model cumulative long-term outcomes in an environment, avoiding autoregressive generation while remaining transferable.
Successor Features Successor features achieve generalization across rewards by modeling the accumulation of features (as opposed to rewards in model-free RL) 2 With the assumption that rewards are linear in features, policy evaluation under new rewards reduces to a linear regression problem. A key limitation of successor features is their inherent policy dependence, they defined as the accumulated features when acting according to particular policy. This makes extracting optimal policies for new tasks challenging.
To circumvent this policy dependence, generalized policy improvement 3 maintains a discrete set of policies and selects the highest valued one to execute at test time, limiting the space of available policies for new tasks. Universal SF [ and Forward- -arkkadd Representations [54] 55 randomly sample reward weights and jointly learn successor features and policies conditioned on The challenge lies in achieving coverage over the space of all possible policies through sampling of z, resulting in potential distribution shifts for new problems. RaMP (9) learns successor feature predictor conditioned on an initial state and a sequence of actions. Transfer requires planning by sampling actions sequences, which becomes quickly intractable over horizon. In contrast, DiSPOs avoid conditioning on any explicit policy representation by modeling the distribution of all possible outcomes represented in a dataset, and then selecting actions corresponding to the most desirable long-term outcome.
Distributional Successor Measure (DSM) [59] is concurrent work that learns distribution over successor representations using tools from distributional RL [4) Importantly, DSM models the distributional successor measure of particular policy, where the stochasticity stems from the policy and the dynamics. This makes it suitable for robust policy evaluation but not for transferring to arbitrary downstream tasks. In contrast, DiSPOs model the distribution of successor feature outcomes in the dataset (i.c.. the behavior policy), where the distribution stems from the range of meaningfully distinct long-term outcomes. This type of modeling allows DiSPOs extract optimal behavior for arbitrary downstream tasks, while DSMs suffer from the same policy dependence that standard successor feature- based methods do.
3 Preliminaries
We adopt the standard Markov Decision Process (MDP) notation and formalism 24 for an MDP M = (A, r, 7; T Pa), but restrict our consideration to the class of deterministic MDPs. While this does not encompass every environment, does capture a significant set of problems of practical interest. Hereafter, we refer to deterministic MDP and a task interchangeably. In our setting, we consider transfer across different tasks that always share the same action space A, state space S, and transition dynamics T S x A > Theh difference between tasks only lies in having different state-dependent Markovian reward functions S + [0, 1).
For simplicity, we also T(s, a) to denote the next state.
3

---

Value Functions and Successor Features Let R R= r(s:) denote the cumulative reward for trajectory {s,, a d= One can then define the state value function under policy as V := B,,T [ 81 = s). and the state-action value function as Q" (s, a) == E [R 81 8, dn = aj The value function admits a temporal structure that allows it to be estimated using dynamic programming. which iteratively applies the Bellman operator until fixed point is reached V (s) = r(s)
y 1xx (V2 si = s). While these Bellman updates are in the tabular setting, equivalent function approximator variants (e.g., with neural networks) can be instantiated to minimize Bellman "error" with stochastic optimization techniques
[3 36).
Successor features [2] generalize the notion of value function from task-specific rewards to task- agnostic features. Given state feature function Φ : Rd, the successor feature of policy is defined as (s) = Ex,T [_r=1 n" o(s;) 81 = s| Suppose rewards can be linearly expressed
by the features, i.e. there exists E R" such that R(s)) w ((s), then the value function for the
particular reward can be linearly expressed by the successor feature (s) w To&T (s)).. Hence, given the successor feature v of a policy т, we can immediately compute its value under any reward once the reward weights w are known. Analogous to value functions, successor features also admit a recursive Bellman identity 15t (s) == Φ(x) + B [ (s')], allowing them to estimated using dynamic programming 3]. In this paper, we also refer to the discounted sum of features along trajectory as a successor feature. In this sense, a successor feature represents an outcome that is feasible under the dynamics and can be achieved by some policy.
Diffusion Models DiSPOs rely on expressive generative models to represent the distribution of successor features. Diffusion models 23 47 are class of generative models where data generation formulated as an iterative denoising process. Specifically, DDPM [23] consists of a forward process that iteratively adds Gaussian noise to the data, and corresponding reverse process that iteratively denoises a unit Gaussian to generate samples from the data distribution. The reverse process leverages a neural network estimating the score function of each noised distribution, trained with a denoising score matching objective [49]. In addition, one can sample from the conditional distribution p(z by adding guidance v, log p(y) (z) the score function in each sampling step [13. As we show in Sec.
guided diffusion enables quick selection of optimal outcomes from DiSPOs.
Problem setting We consider a transfer learning scenario with access to an offline dataset D {(s; a,, < of transition tuples collected with some behavior policy Th under dynamics TT The goal is to quickly obtain the optimal policy x² for some downstream task, specified in the form of a reward function or a number of (s, r) samples. While we cannot hope to extrapolate beyond the dataset (as is common across problems in offline RI. (4), we will aim to find the best policy within dataset coverage for the downstream task. This is defined more precisely in Section
Fi
4 Distributional Successor Features for Zero-Shot Policy Optimization
We introduce the framework of DiSPOs as a scalable approach to the transfer problem described in Section 3 with the goal of learning from an unlabeled dataset to quickly adapt to any downstream task specified by a reward function. We start by relating the technical details behind learning DiSPOs in Section 4.1 followed by explaining how DiSPOs can be used for efficient multi-task transfer in Section Finally, we describe a practical instantiation of DiSPOs in Section 4.3
4.1 Learning Distributional Successor Features of the Behavior Policy
To transfer and obtain optimal policies across different reward functions, generalist decision- mecision-making agents must model the future in a way that permits the evaluation of new rewards and new policies. To this end, DiSPOs adopt a technique based on off-policy dynamic programming to directly model the distribution of cumulative future outcomes, without committing to a particular reward function r(-) or policy л. Fig. 2 illustrates the two components in DiSPOs, and we describe each below.
(1) Outcome model: for particular a state feature function ((s), DiSPOs model the distribution of successor features p(v)s)) over all paths that have coverage in the dataset. In deterministic MDPs,
each successor feature to (discounted sum of features 1/ = Ee --{))((s)) can be regarded as an "outcome". When the state features are chosen such that reward for the desired downstream task is a
linear function of features, i..., there exists w R such that r(s)
(b)s) 42 43 66
56 g Bl.
the value of each outcome can be evaluated as w Tzh. That is, knowing effectively transforms the

---

distribution of outcomes (s) into a distribution of task-specific values (sum of rewards) s). Notably, since can be estimated by regressing rewards from features, distributional evaluation on new reward function boils down to simple linear regression.
As in off-policy RI., the outcome distribution p(y))s) in DiSPOs can be learned via approximate dynamic programming update, which is similar to a distributional Bellman update [4]
max E (s.a,s))>D [log po (Ō(() 7//r1ss]
(I)
s.t
Dy ~ po(-1s")
Intuitively, this update suggests that the distribution of successor features po (v)ss) at state maximizes likelihood over current state feature o(s) added to sampled future outcomes 4/.-. This instantiates fixed-point procedure, much like distributional Bellman update. An additional benefit of the dynamic programming procedure is trajectory stitching, where combinations of subtrajectories the dataset will be represented in the outcome distribution.
Modeling the distribution of future outcomes in an environment is useful only when it can be realized in terms of actions that accomplish particular outcomes. To do so, DiSPOs pair the outcome model with a readout policy [a]ss, +) that actualizes desired long-term outcome & into the action to be taken at state s. Along with the outcome model po (p(lss) the readout policy TT (als,, s, v) can be optimized via maximum-likelitne
(2) Readout policy:
estimation:
max R (") [log (als,, vi = <(s) + rxb.))
s.t
Ve ~ po(.ss')
This update states that if an action a at state leads to a next state 8', then a should be taken with high likelihood for outcomes v., which are a combination the current state feature ((x) and future outcomes a. ~ po(-ss").
4n
Sz
Action
The outcome distribution p(*) can be under- stood as natural analogue to value function, but with two crucial differences: (1) it represents the accumulation of not just single reward fune- tion but an arbitrary feature (with rewards being linear in this feature space), and (2) it is not specific to any particular policy but represents the distribution over all cumulative outcomes covered in the dataset. The first point enables transfer across rewards, while the second ables the selection of optimal actions for new rewards rather than being restricted to partic- ular (potentially suboptimal) policy. Together with the readout policy л (a|s, ))) these models satisfy our desiderata for transfer, i.e., that the value for new tasks can be estimated by simple linear regression without requiring autoregres- sive generation, and that optimal actions can be obtained without additional policy optimization.
0g
Ss
Feature (p)
Successor
Feature (P))
Distributional Successor Features for Zero-Shot Policy Optimization
(1) p(P1s))
(2) (TaaaSo
Figure 2: DiSPOs for simple environment. Given state feature function é, DiSPOs learn distribution of all possible long- outcomes (successor features w))
dataset p(b(s) along with readout policy T (als, Φ) that takes an action a to realise y starting at state
4.2 Zero-Shot Policy Optimization with Distributional Successor Features
To synthesize optimal policies for novel downstream reward functions using DiSPOs, two sub- problems must be solved: (1) inferring the suitable linear reward weights w, for a particular reward function from set of (s. tuples and (2) using the inferred w, to select optimal action a" at state s. We discuss each below.
Inferring task-specific weights with linear regression. As noted, for any reward function r(s), once the linear reward weights w, are known (i.e., r(s) = wi ((s)), the distribution of returns in the dataset p(R(s) is known through lincarity. However, in most cases, rewards are not provided in
5

---

functional form, making wy unknown priori. Instead, given dataset of D = {(s, r)} tuples, can be obtained by solving a simple linear regression problem arg min [D] (..)eD I!w (bs) r
Generating task-specific policies via distributional evaluation. Given the inferred w, and the corresponding future return distribution p [[)s) obtained through linear scaling of p(vls), the optimal action can obtained by finding the v/ with the highest possible future return that has sufficient data-support:
arg max
W W
S.t p(6()) 6) 2
where is a hyperparameter ensure sufficient coverage for sh. This suggests that the optimal outcome 66* is the one that provides the highest future sum of rewards wi o= while being valid under the environment dynamies and dataset coverage.
This optimization problem can be solved in number of ways. The most straightforward is via
Step 1: Linena Regression
Step Planning
random shooting [53], which samples a set of sb from p(t/) s) and chooses the with the high- est ur th. Sec. 5 bases our theoretical analysis on this technique. Sec. 4.3 shows that for out- come models instantiated with diffusion models, the optimization problem can be simplified to guided diffusion sampling.
argggung Ti p(M)) =(ols.0")
w = argmin, liv (bG) - r(
Figure 3: Zero-shot policy optimization with DiSPOs. Once DiSPO is learned, the optimal action can be obtained by performing reward regression and searching for the optimal outcome under the dynamics decode via the policy.
Once 2² has been obtained, the action to exe- cute in the environment can be acquired via the readout policy (a(a, ))²). Fig. 3 shows the full policy optimization procedure, and we refer the reader to Appendix. for the pseudocode. As described previously, DiSPOs enable zero-shot transfer to arbitrary new rewards in an environment without accumulating compounding error or requiring expensive test-time policy optimization. In this way, they can be considered a new class of models of transition dynamics that avoids the typical challenges in model-based RL and successor features.
4.3 Practical Instantiation
In this section, we provide a practical instantiation of DiSPOs that is used throughout our experimental evaluation. The first step to instantiate DiSPOs is to choose expressive state feature that linearly expresses a broad class of rewards. We choose the state fcature é to be d-dimensional random Fourier features [52]. Next, the model class must account for the multimodal nature of outcomes and actions since multiple outcomes can follow from a state, and multiple actions can be taken to realize outcome. To this end, we parametrize both the outcome model p(t)ss and the readout policy π (a)s,
8, (b) using a conditional diffusion model [3 47). We then train these models (optimize Equation 2 by denoising score matching, surrogate of maximum likelihood training [48).
Remarkably, when (1bb))) s) is parameterized by a diffusion model, the special structure of the opti- mization problem in Eq. allows a simple variant of guided diffusion [14 25 to be used for policy optimization. In particular, taking the log of both sides of the constraint and recasting the constrained optimization via the penalty method, we get penalized objective L w & + a p(bls(())))) log e). Taking the gradient yields
V VeL(b, wo + a Vi log p(1))))
The expression for V L(b,, a) is simply the score function V logp(bbl)) in standard diffusion training (Section 3. with the línear weights w, added as a guidance term. Plamning then becomes doing stochastic gradient Langevin dynamics [57] to obtain an optimal 3³ sample, using V V L(b, a) as the gradient. Guided diffusion removes the need for sampling a set of particles. As shown in Appendix F. it matches the performance of random shooting while taking significantly less inference time. Appendix B we show that the guided diffusion procedure can alternatively be viewed as taking actions conditioned on a soft optimality variable.

---

Theoretical Analysis of Distributional Successor Features for Zero-Shot Policy Optimization
To provide theoretical understanding of DiSPOs, we conduct an error analysis to connect the error in estimating the ground truth po(e s) to the suboptimality of the DiSPO policy, and then study when the DiSPO policy becomes optimal. We start our analysis conditioning on >> 0 estimation error in the ground truth outcome distribution po(1) s).
Condition 5.1. We say the learnt outcome distribution p is e-good approximation if S, [ 8) Po s) <.
Since DiSPOs capture the outcome distribution of the behavior policy na, need definition to evaluate policy и with respect to Tg.
Definition 5.2. We say a state-action pair (s.a) is (6, n.) -ggoooodd if over the randomness of Tp. Fa 1i=1QTt0s a) Et= 7f r(()) 81 = s] 6. Furthermore, if for all state s, (s, (s)) is (6, mg)- good, then call π (6, -g)oo policy.
We proceed to use Definition to characterize the suboptimality of the DiSPO policy. Let T denote the sampling optimality of the random shooting planner Sec. [4.2 Specifically, we expect sample a top r outcome & from the behavior policy in ( samples, where P [w] vb S E. (r()) S т. The following result characterizes the suboptimality of the DiSPO policy. The proof is deferred to Appendix. A
Theorem 5.3 (main theorem). For any MDP M and e-good outcome distribution P. the policy f given by the random shooting planner with sampling optimality T is (e ( T, T.gg))ooooo policy.
From Theorem 5.3 we can obtain the following suboptimality guarantee in terms of the value function under the Lipschitzness condition. The corollary shows the estimation error in po s) will be amplified by an O
multiplicative factor.
Corollary 5.4. we have X-Lipschitzness near the optimal policy, i.e., Q'(s, a²) as, X6 when (8, a) is (6, B)-good, the suboptimality of outpul policy î is Vg (so) Vo* (so) S ( +)..
Lastly, we extend our main theoretical result to the standard full data coverage condition in the offline
RL literature, where the dataset contains all transitions [51] 60 44 The following theorem states
that DiSPOs can output the optimal policy in this case. The proof is deferred to Appendix A
Theorem 5.5. In deterministic MDPs, when |A x S) < co, and V(s a) E x A, a, T(s, a)) 2 DiSPOs are guaranteed identify an optimal policy.
Experimental Evaluation
In our experimental evaluation, we aim to answer the following research questions. (1) Can DiSPOs transfer across tasks without expensive test-time policy optimization? (2) Can DiSPOs avoid the challenge of compounding eтror present in model-based RL? (3) Can DiSPOs solve tasks with arbitrary rewards beyond goal-reaching problems? ? (4) Can DiSPOs go beyond the offline dataset, and accomplish "trajectory-sticchinng to actualize outcomes that combine different subtrajectories?
We answer these questions through number of experimental results in simulated roboties problems.
We defer detailed descriptions of domains and baselines to Appendix and C as well as detailed
ablative analysis to Appendix B
6.1 Problem Domains and Datasets
Antmaze 15] is a navigation domain that involves controlling a quadruped to reach some designated goal location. Each task corresponds reaching different goal location. We use the D4RL dataset for pretraining and dense rewards described in Appendix D for adaptation.
Franka Kitchen [5] is manipulation domain where the goal is to control Franka arm to interact with appliances in the kitchen. Each task corresponds to interacting with a set of items. We use the D4RL dataset for pretraining and standard sparse rewards for adaptation.
7

---

(1) Animaze
(2) Franka Kitchen
(3) Hopper
Preference Antmaze
(5) Roboverse
5:3
Figure 4: Evaluation domains: (1) D4RL Antmaze [15] (2) Franka Kitchen [1] (3) Hopper (4) Preference-
Based Antmaze with the goal taking a particular path (5) Roboverse [46 robotic manipulation.
Table I: Offline multitask RL on AntMaze and Kitchen. DiSPOs show superior transfer performance (in average episodic return) than successor features, model-based RI, and misspccified goal-conditioned baselines.
DiSPO (Oun)
USFA
FB
RaMP
MOPO
COMBO
GC-IQL
umaze
593
462 #
469
459 +
451
2
574 + 10
571 : 15
umaze-diverse
568
447
474
460 1
467
547 + 11
577
medium-diverse
631
394 = 52
294
266 +
236
418 + 16
403 =
medium-ple
624
370 31
264
271
232
397 12
390
auree-iivers,
359 +
215 = 20
181
132 +
128
244 = 19
226 =
large-play
306
250
165
134 +
128
248
229
kitchen-parce partial
43
4
0 +
8
11 +
kitchen-mixed
46
10 # 10
5 +
0
hopper-forward
566
487
452
470 16
493 114
982 1:1
hopper-backward
367 + is
261 = 68
269 +
220 + 13
596 2111
194 = 74
hopper-stand
800
685 130
670 1120
255
800
600 111
hopper-jump
832 + 22
746 112
726
652 + 28
753
670 10
Hopper 7 9 is locomotion domain that involves controlling a hopper to perform various tasks,
including hopping forward, hopping backward, standing, and jumping. We use the offline dataset. from 19 for pretraining and shaped rewards for adaptation.
Preference Antmaze is variant of D4RL Antmaze [15] where the goal to reach the top right corner starting from the bottom left corner. The two tasks in this environment correspond to the two paths to reach the goal, simulating human preferences. We collect custom dataset and design reward functions for each preference.
Roboverse [46] is tabletop manipulation environment with robotic arm completing multi-step problems. Each task consists two phases, and the offline dataset contains separate trajectories of each phase but not full task completion. sparse reward is assigned each time step of task completion.
6.2 Baseline Comparisons
Successor Features We compare with three methods from the successor feature line of work. USFA [6] overcomes the policy dependence of SF by randomly sampling reward weights and jointly learning a successor feature predictor th and a policy F. conditioned on z. th captures the successor feature of πT, while π. trained to maximize the reward described by z. FB [54] 55) follows the same paradigm but jointly learns a feature network by parameterizing the successor measure as an inner product between a forward and backward representation. RaMP [ removes the policy dependence of SF by predicting cumulative features from an initial state and an open-loop sequence of actions, which be used for planning.
Model- Based RL We compare with two variants of model-based reinforcement learning. MOPO 64 is model-based offline RL method that learns an ensemble of dynamics models and performs actor-critic learning. COMBO (33 introduces pessimism into MOPO by training the policy using a conservative objective [30].

---

Table Evaluation on non-goal-conditione tasks. DiSPOs are able to solve non-goal-condiitoned tasks, taking different paths in preference antmaze (Fig(4) while goal- conditiiod RL cannot optimize for arbi-
Table 3: Evaluation of trajectory stitching ability of DiSPOs. DiSPOs outperform non-stitching baselines, demonstrating their abilitics to recombine outcomes across trajectory segments
irary rewards.
DiSPO (Ours)
RaMP
DT
DiSPO (Our)
COMBO
GC-IQL
PickPlace ClosedDrawer
49
0 =
0
Up Right
139
143
72
40
5
0
0
0 =
142 + 2
136 = 83 # 25
BlockedDrawer
66
7 0 +
Goal-Conditioned RL Gool-conditioned RL enables adaptation to multiple downstream goals However, it is solving more restricted class of problems than RL goals are less expressive than rewards in the same state space. Moreover, standard GCRL is typically trained on the same set of goals in evaluation, granting them privileged information. To account for this, we consider goal-conditioned RL baseline GC-IQL 40 [9 and only train on goals from half the state space to show its fragility goal distributions. We include the original method trained test-time goals in Appendix
6.3 Do DiSPOs enable zero-shot policy optimization across tasks?
We evaluate DiSPOs on transfer problems, where the dynamies shared, but the reward functions vary. We train DiSPOs on the data distributions provided with the D4RL [15] and Hopper datasets. We identify the test-time re- ward by subsampling small number of transi- tions from the offline dataset, relabeling them with the test- time rewards, and performing lin- ear least squares regression. While DiSPOs in principle can identify the task reward from on- line experience, we evaluate in the offline setting to remove the confounding factor of exploration.
DISPOs
F
3:3
00
Figure 5: Transfer across tasks with DiSPOs and COMBO [63 in medium antmaze. Each tile corre- sponds different task, with color of the tile indicat- ing the normalized return. DiSPOs successfully transfer across a majority of tasks, while MBRL (63 struggles on tasks that are further away from the initial location.
Table rreporsss the episodic return on D4RL and Hopper tasks. DiSPOs are able to transfer to new tasks with no additional training, showing significantly higher performance than successor features (mismatch between training and evaluation policy sets), model-based RL (compounding error) and goal-conditioned RL (goal distribution misspecification). Notably, we show in Appendix [E] that DiSPOs are even competitive with goal- conditioned RI. methods trained on test-time goals. The transferability of DiSPOs can also be seen in Fig 5 where we plot the performance of DiSPOs across various tasks (corresponding to different tiles in the maze). We see that DiSPOs have less degradation across tasks than model-based RL [3
Although the DiSPO framework and theoretical results are derived under deterministic MDPs, we emphasize that the D4RL antmaze datasets are collected with action noise, emulating stochastic transitions. These results indicate that DiSPOs are practically applicable to some range of stochastic settings, although we expect it to perform better in purely deterministic settings.
6.4 Can DiSPOs solve tasks with arbitrary rewards?
While methods like goal-conditioned RL 40 are restricted to shortest path goal-reaching problems, DiSPOs are able solve problems with arbitrary reward functions. This is crucial when the reward is not easily reduced to particular "goal". To validate this, we evaluate DiSPOs on tasks that encode nontrivial human preferences in reward function, such as particular path preferences in antmaze. In this case, we have different rewards that guide the agent specifically down the path to the left and the right, as shown Fig As we see in Table DiSPOs and model-based RL obtain policies that respect human preferences and are performant for various rewards. Goal- conditioned algorithms are unable disambiguate preferences and end up with some probability of taking each path.

---

6.5 Do DiSPOs perform trajectory stitching?
The ability to recover optimal behavior by combining suboptimal trajectories, or "trajectory stitching," is crucial to off-policy RL methods as it ensures data efficiency and avoids requirements for exponen- tial data coverage. DiSPOs naturally enables this type of trajectory stitching via the distributional Bellman backup, recovering "best-in-data" policies for downstream tasks. To evaluate the ability of DiSPOs to perform trajectory stitching, we consider the envirouments introduced in [46]. Here, the data only consists of trajectories that complete individual subtasks (e.g, grasping or placing), while the task of interest rewards the completion of both subtasks. Since the goal of this experiment is to evaluate stitching, not transfer, choose the features the task rewards Φ(s) r(s). We find that DiSPOs are able to show non-trivial success rates by stitching together subtrajectories. Since RaMP 9 predicts the summed features from sequence of actions, and the optimal action sequence is not present in the dataset, it fails to solve any task. Likewise, return-conditioned supervised learning methods like Decision Transformer [10] do not stitch together trajectories and fails to learn meaningful behaviors.
7 Discussion
This work introduced Distributional Successor Features for Zero-Shot Policy Optimization (DiSPOs), method for transferable reinforcement learning that does not incur compounding error or test-time policy optimization. By modeling the distribution of all possible future outcomes along with policies to reach them, DiSPOs can quickly provide optimal policies for any reward in a zero-shot manner. We presented an efficient algorithm to learn DiSPOs and demonstrated the benefits of DiSPOs over standard successor features and model-based RL techniques. The limitations our work open future research opportunities. First, DiSPOs require a choice of features (0x) that linearly express the rewards; this assumption may fail, necessitating more expressive feature learning methods. Second, DiSPOs model the behavior distribution of the dataset; hence, policy optimality can be affected by dataset skewness, which motivates the use of more efficient exploration methods for data collection. Finally, the current version of DiSPOs infer the reward from offline state-reward pairs; potential future direction could apply this paradigm to online adaptation, where the reward is inferred from online interactions.
Acknowledgment
CZ is supported by the UW-Amazon fellowship. TH is supported by the NSF GREP under Grant No. DGE 2140004. SSD acknowledges the support of NSF IIS 2110170, NSF DMS 2134106, NSF CCF 2212261, NSF IIS 2143493, NSF CCF 2019844, and NSF IIS 2229881.
References
[1] K. Asadi, D. Misra, and M. L. Littman. Lipschitz continuity in model-basee reinforcement learning. In J. G. Dy and A. Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässsen, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 264-273, PMLR, 2018
[2] A. Barreto, D. Borsa, J. Quan, T. Schaul, D. Silver, M. Hessel, D. Mankowitz, A. Zidek, and R. Munos. Transfer in deep reinforcement learning using successor features and generalised policy improvement. In International Conference on Machine Learning, pages 501-510. PMLR, 2018.
[3] A. Barreto, W. Dabney, R. Munos, J. J. Hunt, T. Schaul, H. P. van Hasselt, and D. Silver. Successor features for transfer in reinforcement learning. Advances in Neural Information Processing Systems, 30, 2017.
[4] M. G. Bellemare, W. Dabney, and R. Munos. distributional perspective on reinforcement learning. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML. 2017, Sydney, NSW, Australia, I1 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 449-458. PMLR, 2017.
10

---

[5] D. Borsa, A. Barreto, J. Quan, D. J. Mankowitz, H. van Hasselt, R. Munos, D. Silver, and T. Schaul. Universal successor features approximators. In International Conference on Learning Representations, 2019.
[6] D. Borsa, A. Barreto, J. Quan, D. Mankowitz, H. van Hasselt, R. Munos, D. Silver, and T. Schaul. Universal successor features approximators. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
[7] G. Brockman, V. Cheung, L. Pettersson, J. Schneider, J. Schulman, J. Tang, and W. Zaremba. Openai gym. CoRR, abs/1606.015400 2016.
[8] S. Casper, X. Davies, C. Shi, T. K. Gilbert, J. Scheurer, J. Rando, R. Freedman, T. Korbak,
D. Lindner, P. Freire, T. Wang. S. Marks, C. Ségerie, M. Carroll, A. Peng, P. J. K. Christof- fersen, M. Damani, S. Slocum, U. Anwar, A. Siththaranjan, M. Nadeau, E. J. Michaud, J. Pfau, D. Krasheninnikov, X. Chen, L. Langosco, P. Hase, E. Biyik, A. D. Dragan, D. Krueger,
D. Sadigh, and D. Hadeield-Menell. Open problems and fundamental limitations of reinforce- ment learning from human feedback. CoRR, 152177
[9] B. Chen, C. Zhu, P. Agrawal. K. Zhang, and A. Gupta. Self-supervised reinforcement learning that transfers using random features. CoRR, abs/2305. 172500 2023.
17250,
[10] L. Chen, Lu, A. Rajeswaran, Lee, A. Grover, M. Laskin, Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In M. Ran- zato, A. Beygelzimer, Y. N. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurlPS 2021, December 6- 14, 2021, virtual, pages 1 15084-1500 15097, 2021.
[11] K. Chua, R. Calandra, R. McAllister and Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. CoRR, abs/1805. 121 2018.
[12] M. P. Deisenroth and C. E. Rasmussen. PILCO: A model-based and data-efficient approach to policy search. In L. Getoor and T. Scheffer, editors, Proceedings of the 28th International Conference on Machine Learning, 2011, Bellevue, Washington, USA, June 28 July 2. 2011, pages 465 472. Omnipress, 2011.
[13] P. Dhariwal and A. Q. Nichol. Diffusion models beat GANs on image synthesis. In A. Beygelz- imer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[14) Dhariwal and A. Q. Nichol. Diffusion models beat gans on image synthesis. In M. Ranzato, A. Beygelzimer, Y. N. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 8780-8794, 2021.
[15] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4RL: Datasets for deep data-driven reinforcement learning. https:/arrsi.orraab22200...2222 2020.
[16] J. Fu, A. Kumar, O. Nachum, G. Tuckeer, and S. Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR, 0et/22004.02199 2020.
[17] D. Ghosh, A. Gupta, A. Reddy, J. Fu, C. M. Devin, B. Eysenbach, and S. Levine. Learning to reach goals via iterated supervised learning. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[18] A. Gupta, J. Yu, T. Z. Zhao, V. Kumar, A. Rovinsky, K. Xu, T. Devlin, and S. Levine. Reset-free reinforcement learning via multi-task learning: Learning dexterous manipulation behaviors without human intervention. arXiv preprint arXiv:2104.11203 2021.
[19] Haarnoja, A. Zhou, P. Abbeel, and Levine. Soft actor-critie: Off-policy maximum entropy deep reinforeement learning with stochastic actor. arXiv preprint arXiv:1801.0122 2018.
11

---

[20] D. Hafner, T. P. Lillicrap, J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. In 8th International Conference on Learning Representations. ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
[21] D. Hafner, T. P. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. Open 2021.
[22] D. Han, B. Mulyana, V. Stankovic, and S. Cheng. A survey on deep reinforcement learning algorithms for robotic manipulation. Sensors, 23(7):3762, 2023.
[23] Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models, 2020.
[24] R. A. Howard. Dynamic programming and markov processes. John Wiley, 1960.
[25] M. Janner, Y. Du, J. B. Tenenbaum, and S. Levine. Planning with diffusion for flexible behavior synthesis. K. Chaudhuri, S. Jegelka, L.. Song, C. Szepesvári, G. Niu, and S. Sabato, editors, International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 9902-9915. PMLR, 2022.
[26] M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-based policy optimization. In H. M. Wallach, H. Larochelle, A. Beygelzimer, F. d'AlcééBuue, Buc, E. B. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8- 8-14, 2019, Vancouver, BC, Canada, pages 12498-12599, 2019.
[27] M. Janner, I. Mordatch, and S. Levine. Gamma-modelss Generative temporal difference learning for infinite-horizon prediction. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[28] J. Kober, J. A. Bagnell, and J. Peters. Reinforcement learning in robotics: A survey. Int. J. Robotics Res., 32(11):
2013.
[29] I. Kostrikov, A. Nair, and Levine. Offline reinforcement learning with implicit q-learning, In International Conference on Learning Representations, 2022.
[30] A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. CoRR, abs/2006.04777 2020.
[31] N. O. Lambert, K. S. J. Pister, and R. Calandra. Investigating compounding prediction errors in learned dynamics models. CORR, Abs/2203.09637. 2022.
[32] N. O. Lambert, A. Wilcox, H. Zhang, K. S. J. Pister, and R. Calandra. Learning accurate long-term dynamics for model-based reinforcement learning. In 202/ 60th IEEE Conference on Decision and Control (CDC), Austin, TX, USA, December 14-17, 2021, pages 2880-2887 IEEE, 2021.
[33] S. Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. CoRR, abs/1805 50099, 2018
[34] S. Levine, A. Kumar, G. Tucker, and J. Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. CoRR, abs/2005.01 2020.
[35] I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[36] V. Mnih, A. P. Badia, M. Mirza, A. Graves, T. P. Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In M. Balean and K. Q Weinberger, editors, Proceedings of the 33nd International Conference on Ma- chine Learning, ICMI. 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 1928-1937 JMg,. 2016.
12

---

[37] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. A. Ried- miller. Playing atari with deep reinforcement learning. CoRR, abs/1312.5600 2013.
[38] A. Nagabandi, G. Kahn, R. Fearing, and Levine. Neural network dynamies for model- based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE International Conference on Roboties and Automation, ICRA 2018, Brisbane, Australia, May 21-25, 2018, pages 7539-7566. IEEE, 2018
[39] A. Nagabandi, K. Konolige, S. Levine, and V. Kumar. Deep dynamics models for learning dexterous manipulation. In L. P. Kaelbling, D. Kragic, and K. Sugiura, editors, 3rd Annual Conference on Robot Learning, CoRL 2019, Osaka, Japan, October 30 November 1, 2019, Proceedings, volume 100 of Proceedings of Machine Leurning Research, pages 1101-1112. PMLR, 2019.
[40] S. Park, D. Ghosh, B. Eysenbach, and S. Levine. Hiql: Offline goal-conditioned rl with latent states as actions. Advances in Neural Information Processing Systems, 2023,
[41] E. Perez, F. Strub, H. de Vries, V. Dumoulin, and A. C. Courville. Film: Visual reasoning with general conditioning layer. In AAAI, 2018.
[42] A. Rahimi and B. Recht. Random features for large-scale kemnel machines. Advances in Neural Information Processing Systems, 20, 2007.
[43] A. Rahimi and B. Recht. Weighted sums of random kitchen sinks: Replacing minimization with randomization in learning, Advances in Neural Information Processing Systems, 21, 2008.
[44] T. Ren, J. Li, B. Dai, S. S. Du, and S. Sanghavi. Nearly horizon-free offline reinforcement learning. Advances in neural information processing systems, 14:15111-5 2021.
[45] O. Rybkin, C. Zhu, A. Nagabandi, K. Daniilidis, I. Mordatch, and Levine. Model-based reinforcement learning via latent-space collocation. In M. Meila and T. Zhang, editors, Proceed- ings of the 38th International Conference on Machine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 9190-9201. PMLR, 2021.
[46] A. Singh, A. Yu, J. Yang, J. Zhang, A. Kumar, and Levine. Cog: Connecting new skills to past experience with offline reinforcement learning. Preprint arXiv:2010. 2020.
[47] J. Song, C. Meng, and S. Ermon. Denoising diffusion implicit models, 2022.
[48] Y. Song, C. Durkan, I. Murray, and S. Ermon. Maximum likelihood training of score-based diffiusion models. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021.
[49] Y. Song, J. Sohl-Dickstein, D. P. Kingma, A. Kumar, S. Ermon, and B. Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations, 2021.
[50] R. S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bull., 164):160- 1991.
[51] C. Szepesvári and R. Munos. Finite time bounds for sampling based fitted value iteration. In Proceedings of the 22nd international conference Machine learning, pages 880-887, 2005.
[52] M. Tancik, P. P. Srinivasan, B. Mildenhall, S. Fridovich-Keil, N. Raghavan, U. Singhal, R. Ra- mamoorthi, J. T. Barron, and R. Ng. Fourier features let networks learn high frequency functions in low dimensional domains. NewrIPS, 2020.
[53] R. Tedrake. Underactuated Robotics. 2023.
[54] A. Touati and Y. Ollivier. Learning one representation to optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23, 2021.
[55] A. Touati, J. Rapin, and Y. Ollivier. Does zero-shot reinforcement learning exist? In The Eleventh International Conference on Learning Representations, 2023.
13

---

[56] A. Wagenmaker, G. Shi, and K. Jamieson. Optimal exploration for model- based RL in nonlinear systems. CoRR, Abs/2306.092111 2023.
[57] M. Welling and Y. W. Teh. Bayesian learning via stochastic gradient langevin dynamics. In L. Getoor and T. Scheffer, editors, Proceedings the 28th International Conference on Machine Learning, ICML 2011, Bellevue, Washington, USA, June 28 July 2, 2011, pages 681-688. Omnipress, 2011.
[58] G. Williams, N. Wagener, B. Goldfain, P. Drews, J. M. Rehg, B. Boots, and E. A. Theodorou. Information theoretic MPC for model-based reinforcement learning. In 2017 IEEE International Conference Robotics and Automation, ICRA 2017, Singapore, Singapore, May 29 June 2017, pages 1714-1721. IEEE, 2017.
[59] H. Wiltzer, J. Farebrother, A. Gretton, Y. Tang. A. Barreto, W. Dabney, M. G. Bellemare, and M. Rowland. distributional analogue to the successor representation, 2024.
[60] Xie and Jiang. Batch value-function approximation with only realizability. International Conference Machine Learning, pages 1404-11413, PMLR, 2021.
[61] Y. Xu, J. Parker- A. Pacchiano. P. J. Ball, O. Rybkin, S. Roberts, T. Rocktäschel, and E. Grefenstette. Learning general world models in a handful of reward-free deployments. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
[62] K. J. Young, A. Ramesh, L. Kirsch, and J. Schmidhuber. The benefits of model-based general- ization in reinforcement learning. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, International Conference on Machine Learning. ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 40254 40276. PMLR, 2023.
[63] T. Yu, A. Kumar, R. Rafailov, A. Rajeswaran, S. Levine, and C. Finn. Combo: Conservative offline model-based policy optimization. M. Ranzato, A. Beygelzimer, Y. Dauphin, Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 28954 28967. Curran Associates, Inc., 2021.
[64] T. Yu, G. Thomas, L. Yu, S. Ermon, J. Zou, S. Levine, C. Finn, and T. Ma. Mopo: Model-based offline policy optimization. Preprint arXiv:2005. 13239, 2020.
[65] M. R. Zhang, Paine, O. Nachum, C. Paduraru, G. Tucker, Z. Wang, and M. Norouzi. Autore- gressive dynamics models for offline policy evaluation and optimization. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[66] T. Zhang, T. Ren, M. Yang, J. Gonzalez, D. Schuurmans, and B. Dai. Making linear mdps practi- cal via contrastive representation learning. K. Chaudhuri, S. Jegelka, L.. Song, C. Szepesvári, G. Niu, and S. Sabato, editors, International Conference Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 26447-26466. PMLR, 2022.
14

---

Supplementary Materials for "Distributional Successor Features Enable Zero-Shot Policy Optimization"
Missing Proofs
We provide the complete proofs of theorems and corollaries stated in Sec. following sections, we use 1ş to denote the indicator of event &
5
Throughout the
A.1 Formal Statement and Proof of Theorem 5.3
To state Theorem 5.3 rigorously, we introduce the basic setting here. Without loss of generality, let the feature at time step ((s;) E [0, 1 > and the outcome to E -E=o' (';(s;) [0, 1] Moreover, we have a readout policy æ s.t. ~ (b) always leads to a successor state s' s.t. P( (v6 o(s) s') 0.
We simplify the policy optimization phase of DiSPOs into the following form: for given reward weight Wy, each time step, we have
1. Infer optimal outcome a A(s, P) through random shooting;
2. Get corresponding action from मt à f(s, v²
where the random shooting oracle A with sampling optimality t satisfies
f 1 s) S T}.
oh A(s 2 min(R)
This intuitively means that there are at most probability E [0, 1] of the behavior policy achieving higher reward. We proceed to prove Theorem 5.3
Proof. It suffices to prove that Vs, (s, +()) is at least + +,,))) For simplicity, denote R :== min R
sldb S 7.. Then
P m ((s))) 2
/ 1 ( po( (bb s)de
/
s)d
/ 1 s)dt
/ 11 (po(s) s) p(6) s)] dib
$ T / /1ugw22)
<T+6
15

---

A.2 Proof of Corollary 5.4
Proof. Intuitively, policy fall behind by at most A(T + e) each time step: Vsi S.
(s))) - (s1)
= -"(,s,, (8),a)
= - -
(1, Q* (s)
S <(T + (s),
(5),
= + + / (52) (32)
s +(+ c) +
E, V' (8a) - (sa))
7(T +c)
0m0
X
(r + c).
A.3 Proof of Theorem 5.5
Proof. Under the full coverage condition, the sampling optimality can be set to be zero for concrete actions, and the condition p(tba) s) > e becomes p(vb(s) >> Then we see that planning only focuses on the supporting set of p: p( s) if and only (s. 1b) D for some time during execution, which equals po( s) = 0. This indicates that supp(p) supp(pu). Therefore the conclusion follows.
B Alternative Derivation of Guided Diffusion Sampling
In this section, we derive the guided diffusion sampling from the perspective of control as inference B3.. To start, define the trajectory-level optimality variable O Bernoulli variable taking the value of with probability exp(R(r)) and 0 otherwise, where R(T) Σt 7'r(s;) Rmax- Note we subtract the max discounted return Rmax to make the density a valid probability distribution. Planning can be cast as an inference problem where the goal is to sample ui* ~ p(pp((00) O). By Bayes
rule, we have
OO#1O) p(O(v)p(()
Taking the gradient of the log of both sides, we get
V log p(ú O) = V logp(Olp)) (1))) + V log
Vo log exp(w' T +)) + V log
Vew To V V (1))(
= V log
This implies that we sample from O() by adding the regression weights the score at each timestep, yielding the same guided diffusion form as in Sec. 4.2
C Implementation Details
C.1 Model architecture
We parameterize the random Fourier features using a randomly initialized 2-layer ML.P with 2048 units in each hidden layer, followed by sine and cosine activations. For d-dimensional feature,
16

---

the network's 's output dimension is [d/2] and the final feature is concatenation of sine and cosine activated outputs. We set d 128 for all of our experiments.
We implement the outcome model and policy using conditional DDIMs [47] The noise prediction network is implemented
1-D Unet with down dimensions [256, 512, 1024). Each layer is modulated using FiLM 41 to support conditioning.
C.2 Training details
We train our models on the offline dataset for 100,000 gradient steps using the Adam optimizer [35] with batch size 2048. The learning rate for the outcome model and the policy are set to 3e and adjusted according to cosine learning rate schedule with 500 warmup steps. We train the diffusion noise prediction network with 1000 diffusion timesteps and sample using the DDIM +7 sampler with 50 timesteps. We sample from an exponential moving average model with decay rate 0.995. The same set of training hyperparameters is shared across all environments.
To transfer to downstream task, we randomly sample 10000 transitions from the dataset, relabel them with the test time rewards, and perform linear least squares regression. We use the guided diffusion planner for Antmaze (a = 0.5), Franka Kitchen (a 0.01), and Roboverse (a = 0.05) experiments. For Hopper, we use the random shooting planner with 1000 particles. We found planning with guided diffusion to be sensitive to the guidance coefficient. Hence, for new environments, we suggest using the random shooting planner to get a baseline performance and then tuning the guided diffusion coefficient to acclerate inference. We run each experiment with 6 random seeds and report mean and standard deviation in the tables. Each experiment (pretraining + adaptation) takes hours on single Nvidia L40 GPU.
C.3 Baselines
Successor Features We use the implementation of Universal Successor Features Approximators Representation [54] from the code release of (55). We choose random
[5] and Fourier features for universal SF as it performs best across the evaluation suite. To ensure fairness of comparison, we set the feature dimension to be 128 for both Universal SF and FB. Both methods are pretrained for million gradient steps and adapted using the same reward estimation method as DiSPOs.
We adapt the original RaMP implementation [9] and convert it into an offline method. RaMP originally consists of an offline training and an online adaptation stage, where online adaptation alternates between data collection and linear regression. We instead adapt by subsampling 10000 transitions from the the offline dataset and relabeling them with the test-time reward function, thus removing the exploration challenge. We use an MPC horizon of 15 for all experiments.
RaMP
Model-based RL We the original implementations of MOPO 64 and COMBO 663 in evaluations. We pretrain only the trainsition model on the offline transition datasets. To transfer to downstream tasks, freeze the transition model, train the reward model on state reward pairs, and optimize the policy using model-based rollouts. We set the model rollout length for both methods to 5 and the CQL coefficient to be 0.5 for COMBO.
Goal-conditioned RL We use the GC-IQL baseline from 40). tion, we modify the sampling distribution to only sample from half of the goal space excluding the test-time goal location.
To remove the privileged informa-
D Environment Details
Antmaze [6 is navigation domain that involves controlling an 8-DoF quadruped robot to reach some designated goal location in maze. Each task corresponds to reaching different goal location. We use the standard D4RL offline dataset for pretraining. For downstream task adaptation, we replace the standard sparse reward (s g) with a dense reward exp(-l|s gll//20) mitigate the challenge of sparse reward in long- horizoorizon problems.
17

---

122-
124
124
38
23
14
14
Figure 6: Data distribution and reward for Antmaze Preference environments. The left figure illustrates preference for taking the vertical path, and the right figure illustrates the preference for taking the horizontal path.
Franka Kitchen [8] is a manipulation domain where the goal is to control Franka arm to interact with appliances in the kitchen. Each task corresponds interacting with set of items in no particular order. We use the standard D4RL offline dataset for pretraining. For downstream task adaptation, the Markovian sparse rewards, where each timestep the robot gets reward equal to the number completed tasks. We report the number of tasks completed throughout the entire episode in Table
locomotion domain where the task to control hopper to perform various tasks, including hopping forward, hopping backward, standing, and jumping. We design shaped reward for each task. We pretain on the offline dataset from D. collected by taking the replay buffers of expert SAC [9 agents trained to solve each task.
Hopper Hopper
Preference Antmaze is variant of D4RL Antmaze 115 where the goal is reach the top right cell from the bottom left cell in a custom maze shown in Fig. The two tasks in this environment are the two paths to reaching the goal, simulating different human preferences. To construct the dataset, we collect I million transitions using the D4RL waypoint controller. For each preference, we design reward function that encourages the agent take one path and the other. Fig. 6 visualizes the dataset and the reward function for each preference.
Roboverse [46] is tabletop manipulation environment consisting of WidowX arm aiming to complete multi-step problems. Each task consists of two phases, and the offline dataset contains separate trajectories for each phasebut not full task completion. We the standard sparse reward, assigning a reward of for each timestep the task is completed.
E
Additional Experiments
To understand the impact of various design decisions on the performance of DiSPOs, conducted systematic ablations on the various components, using the Antmaze medium-diverse as a test bed.
E.1 Full D4RL Results
Table. 4 displays the full D4RL results with oracle goal-conditioned baseline labeled GC-Oracle. The baseline is trained on a goal distribution covering the testing-time goals, granting privileged information. Despite this, DiSPOs competitive with the oracle in most domains.
E.2 Ablation of Planning Method
We compare the guided diffusion planner with the random shooting planner described in Sec. 4.2 As shown in Table 5 the guided diffusion planner achieves comparable performance to random shooting with 1000 particles while taking significantly less wall-clock time. While we can decrease the number of samples the random shooting planner improve planning speed, this comes the cost of optimality
18

---

Table 4: Full offline multitask RI AntMaze and Kitchen. DiSPOs show superior transfer performance (in average episodic return) than successor features, model-based RL, misspecified goal-conditioned baselines. while being competitive with an oracle using privileged information.
DiSPO (Ours)
USFA
FB
RaMP
MОPO
COMBO
GC-IQL
GC-Oracle
umaze
593 16
462 +
469 # 12
459
451 +
574 +
571 # 15
623 =
umaze-diverse
568
447
474
460
467
547
577 #
576 =
medium-diverse
631 = 67
394
294 +
266
236
418 +
403 + 10
659 i 44
medium-play
624
370
264
271
232
397
390 = 33
673 = 45
large-diverse
359 89
215
181
132
128
244
226 +
493 =
large-play
306 18
250
165 # 12
134
128
248
229 +
533 #
kitchen-partia
=
0
5 #
8
11
= 23
kitchen-mixed
46
10
0
0
43
hopper-forward
566 =
487 1:10
452
470
493 1114
982 1157
hopper-backward
367 15
261
269 77
220
596 1:111
194 +: 74
hopper-stand
800 = 685 1:30
670 =
255
800
600 1111
hopper-jump
832 22
746 1112
726 35
652 28
753
670 1
Table 5: Ablation of planning method. Wall time measured over 1000 planning steps.
Return t
Wall time (s) |
DiSPO (Ours)
631 + 67
42.9
Random shooting @ 1000
650
94.8
Random shooting @ 100
619 = 90
58.6
Random shooting @ 10
513 = 52
55.5
E.3 Ablation of Feature Dimension and Type
To understand the importance of feature dimension and type, we compare variants of our method that use liver-dimensionll random Fourier features, vanilla random features, and the two top-performing pretrained features from (55). From Table. 6 we observe that as feature dimension decreases, their expressivity diminishes, resulting in lower performance. We found random features to perform much worse than random Fourier features. Interestingly, pretrained features with dynamics prediction and graph Laplacian objectives also achieve lower returns than random Fourier features. We hypothesize these pretrained features overfit to the training objective and are less expressive than random Fourier features
E.4 Ablation of Dataset Coverage
We investigate the effect of dataset coverage on the performance our method. We compare DiSPOs trained on the full D4RL dataset against two variants, where randomly subsample half of the transitions, and the other where we adversarially remove the transitions from the half of the state space containing the test-time goal. As shown in Table 7 the performance of DiSPO drops dataset coverage degrades.
E.5
Ablation of Planning Horizon
In Fig. g we ablate the effective planning horizon of DiSPOs by controlling the discount factor y. We found that for shorter horizon tasks (umaze-diverse)) the optimality of planned trajectories improves the planning horizon decreases. However, for hong-horizon tasks (medium-diverse), reducing the planning horizon too much incurs cost on global optimality.
E.6 Nonparametric Baseline
To confirm the intuition of our method, we implemented a nonparametric baseline that constructs an empirical estimate of the outcome distribution. First, calculate the empirical discounted sum of
19

---

Table Ablation feature dimension and type.
Table 7: Ablation dataset coverage.
Return f
Return
DiSPO (Ours)
631 i
Full dataset Random Subsampling
631 = 67
Random Fourier (64-dim)
561 i 45
459 57
Random Fourier (32-dim)
295 = 30
Adversarial Subsampling
390
Random Fourier (16-dim)
307 38
Random
382 +: 43
Table 8: Ablation of planning horizon.
Forward dynamics
402 36
umaze-diverse medium-divers.
Laplacian
376 4: 33
0.99
587
650
7 0.95
597 + 16
682 : 20
7 0.9
601
553 36
Table 9: Comparison to nonparametric baseline that takes the top-valued action among the & nearest neighbors of state.
Return 1
DiSPO (Ours)
631 #: 67
Nonparametric k = 10
308 +: 20
Nonparametrie k 100
299 4
Nonparametric ic 1000
287 + 10
features along trajectories in the dataset. Given the of empirical (s, a, v/) pairs and the downstream reward weight w, we can then select the optimal action at state through the following steps: (1) query the nearest neighbors of 8, (2) evaluate their corresponding values w T vb, (3) take the action of the top-valued neighbor.
We found this baseline to perform surprisingly well on sawiaaze-mmiium-iiiieeee not achieve the performance of DiSPOs, it outperforms RaMP and MOPO. This result confirms the intuition behind DiSPOs, which involves selecting the optimal outcome under dataset coverage and taking an action to realize it. We attribute the performance gap between this baseline and DiSPOs to their trajectory stitching ability (acquired via dynamic programming), infinite horizon modeling, and neural network generalization.
While it does
E.7 Visualization of Stitching
We visualize the stitched trajectories for roboverse environments in Fig. The dataset only contains trajectories from each phase separately, but DiSPOs can generate full trajectories by stitching the subtrajectories.
Phase
Phase 2
Figure 7: Visualization of stitched trajectories for roboverse PickPlace and BrowkedDrawer.
20

---

Algorithm Pseudocode
Algorithm DiSPO Training
1: Given transition dataset D, feature function o6()
2: Initialize po (v(()s), (a(als 16).
3: while not converged do
4: Draw B transition tuples {st, at, s D.
5: Sample successor features for next states 1/1 ~ po (1 (/ == ....N.
6: Construct target successor feature v/ o(s) + V/
7:
// & model learning
8: Update feature distribution: < arg max, log pe u,, mir ss)).
9:
// Policy extraction
10: Update policy: p + arg max, log (alass; ).
end while
Algorithm 2 DiSPO Offline Adaptation
Given transition dataset D, feature function o(-), reward function r(s).
2: Relabel offline dataset 'D with reward function.
3: Initialize regression weights w.
4:
Fit to 'D using linear regression arg min EDll w (T((()))))) - rr()) 11
Algorithm 3 DiSPO Online Adaptation
1: Given pe (1))s), x (als, v)), feature function d(.).
2: Prefill online buffer Dour with random exploration policy.
3: Initialize regression weights w.
4: for time steps 1, T do do
5: Fit w using linear regression w = arg min Eouil w ((8) - rs)) 113).
6: Infer optimal 1 = arg max w vb s.t. (a(()))) e.
7: Sample optimal action a" (als, vd)).
8: Execute action in the environment and add transition to buffer.
9: end for
Algorithm DiSPO Inference (Random Shooting)
Given (vbls,, (apa)s v)), regression weight w, current state
2: Sample N outcomes ( ~ po (w)ss.
3: Compute corresponding values (v: J. where vi w TTTz the
4: Take optimal cumulant s/* v/s, where arg max; (v:
5: Sample optimal action a" ~ Tals, (b).
Algorithm 5 DiSPO Inference (Guided diffusion)
1: Given diffusion model (1())))) (a(a,,, v), regression weight w, current state 8, guidance coefficient 3.
2: Initialize outcome 1 from prior.
1...7 do 4: Compute noise at timestep ε = e t, s).
3: for diffusion timestep
5: Update noise é'
BV1 T.w.
6: Sample next timestep action riz+1 using c'.
7: end for
8: Sample optimal action a" ~ (a,ss v/)).
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction clearly state the contributions and scope of the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are addressed in the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
22

---

The paper discusses the assumptions in Section 5 and provide full proofs in Appendix A
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: The paper disclose the experiment details, including environments, datasets, and model hyperparameters in the Appendix C These can be used to fully reproduce the experiments the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
23

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We include the code in supplemental materials and will release the code if the paper is accepted.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
Policy)lllyy for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips.ccpublle pwwppppguuuiiieeeep httii///uiidessoeeuuuiiii SooodeSubiissi.e ssiiiioicccy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The training and test details are specified in Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: The paper includes error statistics for all reported results.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi-
dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: The paper discloses information about the compute resources used for experi- ments in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't It make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
The paper conforms to the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction..
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
The paper is algorithmic in nature and involves experiments in simulated environments. Hence it does not have an immediate societal impact in its current form. However, may have an impact when deployed in real-world settings where modeling is important, such as robotics or autonomous driving. In these cases, natural challenge with GOMs that the interpretability that comes with autoregressive generative models is naturally lost, since only cumulative features are modeled. For future deployment of this work, we must carefully consider interpretability and think about how we can train generative heads interpret the various outcomes that are modeled by GOMs.
Justification:
25

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not involve the release of models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The paper credits the owners of codebases and datasets that are used in the experiments.
Justification:
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
26

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thcode com comddatasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not introduce new assets.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve crowdsoureing or experiments with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve human subjects and hence does not require IRB approval.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
27

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
28


Reviews:
[{"review_id": "uB1LMDFDZq", "reviewer": "Reviewer_7fQM", "summary": "The paper proposes a method based on successor features for modeling possible long-term outcomes in the environment based on data, together with a learned policy to achieve those outcomes. After training from offline data, the model can produce a policy for a given new reward function without any additional interaction with the environment.", "questions": "Can a model produce any behavior that cannot be done by explicit stitching of trajectory chunks from the training data? Does one of the existing evaluations showcase this?\n\nIs there a way, based on current experimental results, to reason about the efficiency of finding the desired outcome as the domain becomes more and more complex? Is this part expected to become the bottleneck in those cases?", "limitations": "It would be good to have a more clear comparison of pros on cons compared to baseline approaches the method is tested against, in particular related to some of the points raised above.\n\nAdditional evaluations would be valuable to be able to better judge the performance of the approach. Possibly, some of the evaluation environments from the RaMP paper could be used. Even with the proposed approach being fully offline it would be interesting to see how long it takes for methods evaluated there to catch up or exceed its performance. Most of the relevant baselines are already present in those evaluations -- the authors could possibly just run their approach and report those results.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The general idea seems promising. Modeling long-term outcomes rather than single-step dynamics, like in regular model-based approaches, has the potential to help with the problem of error accumulation. This seems to be reflected in the evaluations performed.\n\nThe approach seems to outperform RaMP (although a modified version, made to be fully offline), another successor-feature-based method aiming to solve a similar issue of not being tied to a specific policy and being able to efficiently find policies optimizing for novel rewards.", "weaknesses": "1. The approach is explicitly modeling the distribution of outcomes in the training data. While it should have the ability to reach known outcomes from different initial configurations by performing trajectory stitching, it is unclear whether it could achieve any outcome not explicitly seen in the training data, even if it could be a simple combination of known ones. In contrast, model-based methods have the potential to allow for this, by reaching novel states by rolling out learned local models on novel action trajectories.\n\n2. Another concern is related to finding the outcome that maximizes the given reward. Sampling in the space of all possible outcomes will not be scalable to more complicated environments, therefore it seems critical for the proposed conditional diffusion model to be able to find the desired outcome efficiently. It is difficult to judge based on the presented results that this is the case. In the ablation comparing the proposed method with random shooting, random shooting with 1000 samples outperforms the proposed method slightly while only taking twice as long to execute. It seems that for effectively finding the desired outcome in more complex environments this ratio would need to be much bigger, as we know that sampling will not be close to efficient enough.\n\n3. The approach assumes access to the reward function in a functional form (or at least the value of it across the entire training dataset). It is unclear how realistic this assumption is. Other similar methods (like RaMP) use only the value of the new reward gotten from online interactions with the environment at test time.\n\nThe evaluations present might be insufficient to fully judge the quality of the approach:\n\n4. Goal-Conditioned RL is trained on goals from only one half of the state space, while the proposed method is given data from the entire state space. This seems like inherently an unfair comparison. It would only be fair to give the same distribution to the proposed method for comparison.\n\n5. Both the non-goal-conditioned evaluation and the trajectory stitching evaluation show simply that the method is capable of something in a most basic setup, beating only methods that by design are not capable of doing what is being evaluated. There should be more extensive evaluations to test how well they perform each of these functions, compared with other methods that might be plausibly used in such a case.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and for finding our approach “promising.” We address your comments below.\n\n> (W1) Model-based RL can generate novel action trajectories whereas GOM is constrained to the dataset distribution.\n\nWhile model-based RL can generate novel trajectories when queried with out-of-distribution actions, these trajectories are not guaranteed to be accurate since the model has never seen these transitions during training. Optimizing the policy under hallucinated trajectories can lead to model exploitation, resulting in suboptimal behavior when the policy goes out-of-distribution. In general, the best we can do in offline RL (either model-based or model-free) is to find the best in-distribution trajectory [1]. There can be in-distribution interpolation to new states, but not extrapolation to out-of-distribution states. In this sense, GOM and model-based RL have equal capabilities.\n\n> (W2, Q2) Concern about planning efficiency.\n\nPlanning efficiency is a common concern shared by model-based planning methods, as one needs to search over actions over a long horizon in each planning step. In practice, planning with models is still reasonable because we typically don’t have to plan over the entire task horizon, just a shorter horizon (e.g. by dropping the discount factor) followed by replanning (as in model predictive control). This type of MPC reduces planning complexity while maintaining overall performance. Empirically, we experiment with planning over shorter horizons by decreasing the discount factor. As shown in Table 4 in the attached pdf, for short-horizon tasks such as antmaze-umaze, reducing the planning horizon indeed leads to better results when fixing the planning budget. On the other hand, for longer-horizon problems such as antmaze-medium, as one would expect, reducing planning horizon comes at the cost of global optimality.\n\n> (W3) The approach assumes access to the reward function in a functional form (or at least the value of it across the entire training dataset). It is unclear how realistic this assumption is.\n\nOur method does not inherently require access to the reward function in a functional form. Rather, it only requires a dataset of $(s, r)$ pairs to perform reward regression. This dataset can be collected using an exploration policy, or using the current task policy in an online setting, or by relabeling existing datasets when we have functional rewards. In particular, the reward relabeling scheme described in the paper is equivalent to rolling out the behavior policy to collect a new dataset for each new reward. \n\n> (W4) Goal-conditioned RL unfair comparison.\n\nWe note that hindsight goal-conditioned RL consists of two distributions: the data distribution and the goal relabelling distribution. In our experiments, we remove the test-time goal from the goal-relabeling distribution, so the agent is still trained on all the transition data, but does not see the test-time goal. That said, we acknowledge that the comparison with goal-conditioned RL is rather delicate. With full goal converge, GCRL will be tested on training data, an unfair advantage compared to GOM which does not see test time reward function. On the other hand, covering goals renders GCRL incapable of learning behaviors to reach the goal. We provide results with both misspecified goal distribution and full goal distribution in Table 1 and 4 (Appendix) respectively to give a complete picture. We further emphasize that goal-conditioned RL is not generally applicable to the family of tasks we are considering, as shown in Table 5 of the attached pdf, while GOMs are naturally applicable to any task with Markovian rewards. \n\n> (W5) Non-goal-conditioned and trajectory stitching experiments show capability but not better performance to comparable methods.\n\nThe experiments in Sections 6.4 and 6.5 are designed for analytical purposes to highlight the capability of GOMs. We do provide more experiments comparing the *performance* of GOM in terms of each capability. Specifically,  the kitchen-mixed environment in Table 3 of the paper requires explicitly stitching trajectories of completing different subtasks, as the dataset does not contain full task trajectory. We see that GOM outperforms baselines with trajectory stitching capability. To compare the performance of GOM on challenging non-goal-reaching tasks, we conduct additional experiments on the Hopper environment adapted from RAMP, where the agent is tasked with moving forward, moving backwards, standing, or jumping. As shown in Table 1, our method demonstrates competitive performance compared to the baselines, achieving the highest average return across 4 tasks.\n\n> (Q1) Can a model produce any behavior that cannot be done by explicit stitching of trajectory chunks from the training data? Does one of the existing evaluations showcase this?\n\nStitching is typically achieved by dynamic programming (e.g. Bellman backup). GOM enables stitching via distributional Bellman backup, and hence exhibits the same stitching behavior as standard offline RL methods such as CQL, with the additional capability of transferring to new rewards.\n\n> (Limitations) Additional evaluations would be valuable to be able to better judge the performance of the approach. Possibly, some of the evaluation environments from the RaMP paper could be used.\n\nThanks for suggesting additional evaluation environments. We adapt the Hopper environment from RaMP to the offline setting and compare our methods to representative baselines. We found our method to be competitive, achieving the highest average return across 4 tasks.\n\nReferences:\n\n[1] Sergey Levine, Aviral Kumar, George Tucker, Justin Fu. Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems. ArXiv 2020.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "48mAP4AWni", "reviewer": "Reviewer_AWqu", "summary": "This paper proposes an approach to zero-shot reinforcement learning through learning the distribution of successor features in deterministic MDPs allowing for the efficient computation of approximately optimal policies. The authors perform an empirical evaluation comparing their method to other zero-shot RL, model-based RL, and goal-conditioned RL methods showing gains in various continuous control domains.", "questions": "- You claim that random Fourier features performed best for USFA. Can you provide results for this? I would expect to see a comparison with either the Laplacian features from [1] or the HILP features from [2]. In fact, [2] shows that random features (not fourier) perform poorly compared to better base features for USFA.\n- You chose to set the embedding dimensionality of FB to 128, did you sweep over this value? The embedding dimensionality has a regularizing effect in FB and in many tasks needs to be decreased to obtain good performance.\n- Section 6.4, GOMs aren't the only method that can solve tasks with arbitrary rewards, why wasn't a comparison with USFA / FB performed here?\n- How is the objective function (1) an off-policy update? This is a textbook example of an on-policy update, the target policy is the behavior policy (mixture policy over the dataset) by definition here.\n\n[1] Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? International Conference on Learning Representations (ICLR), 2023.\n\n[2] Park Seohong, Tobias Kreiman, and Sergey Levine. Foundation Policies with Hilbert Representations. International Conference on Machine Learning (ICML), 2024.", "limitations": "I don't believe the authors adequately addressed the limitations of their work. I believe assuming deterministic MDPs is too restrictive in practice. It's unclear how their method performs under any form of environment stochasticity which other methods like USFA, FB, and DSM do account for. At minimum, it would be nice for the authors to discuss this limitation in greater detail providing insight for how we can move beyond this limitation.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The way the authors learn a distribution over successor features is novel, i.e., using a maximum likelihood approach versus prior works that use MMD [1, 2].\n- Using a diffusion model as a generative model of features through a loss employing bootstrapping is unique to this work as well as its application for planning.\n- The paper is generally well-written and easy to read.\n\n\n[1] Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, and Tie-Yan Liu. Distributional Reinforcement Learning for Multi-Dimensional Reward Functions. Neural Information Processing Systems (NeurIPS), 2021.\n\n[2] Harley Wiltzer, Jesse Farebrother, Arthur Gretton, Yunhao Tang, André Barreto, Will Dabney, Marc G. Bellemare, and Mark Rowland. A Distributional Analogue to the Successor Representation. International Conference on Machine Learning (ICML), 2024.", "weaknesses": "I have some serious issues with the framing of this paper. There are many contradictory statements and I believe the paper is a useful contribution but suffers from poor presentation and misleading claims. I'll outline my major concerns:\n\n- I don't think the term generalized occupancy model is appropriate here. First, this isn't an occupancy model, we're modeling successor features not state occupancy as in [1]. Secondly, I don't believe the model is general in any sense of the word. As I'll discuss the method is policy dependent and is currently limited to deterministic MDPs.\n- The paper gives a very confusing and somewhat misleading characterization of the policy-dependent nature of their method. There are many instances of this and I'll try to outline some of them:\n\t- In the introduction, \"Rather than modeling the successor features under a particular policy, GOMs model the entire distribution of successor features under the behavior policy, ...\", isn't the behavior policy not a particular policy?\n\t- \"Importantly, DSM models the distributional successor measure of a particular policy, where the stochasticity stems purely from the policy and the dynamics. This makes it suitable for robust policy evaluation but not for transferring to arbitrary downstream tasks\", If the DSM has these limitations and GOMs don't could you explain the difference between modeling the DSM over the mixture policy induced by a dataset when you have deterministic dynamics? If the answer is there's no difference then I don't think the framing of the DSM vs GOMs is a fair characterization.\n\t- \"A key limitation of successor features is their inherent dependence on a single policy, as they are defined as the accumulated features when acting according to a particular policy. This makes extracting optimal policies for new tasks challenging.\", again, given the problem setting you describe (i.e., fixed dataset, deterministic dynamics) I fail to see how this is a limitation. Learning SFs on this dataset have the same policy dependence as GOMs.\n- The proposed method is closely related to 𝛾-models [2] and geometric horizon models [3]. Under deterministic dynamics, there's not much algorithmic novelty as GOMs reduce to the same cross-entropy TD objective in [2, 3] trained with a diffusion model.\n- I believe the experimental methodology has some room for improvement, especially around the comparison to USFA and FB. See my questions below.\n\n[1] Harley Wiltzer, Jesse Farebrother, Arthur Gretton, Yunhao Tang, André Barreto, Will Dabney, Marc G. Bellemare, Mark Rowland. A Distributional Analogue to the Successor Representation. CoRR abs/2402.08530, (2024)\n\n[2] Michael Janner, Igor Mordatch, and Sergey Levine. Gamma-Models: Generative Temporal Difference Learning for Infinite-Horizon Prediction. Neural Information Processing Systems (NeurIPS), 2020.\n\n[3] Shantanu Thakoor, Mark Rowland, Diana Borsa, Will Dabney, Rémi Munos, and André Barreto. Generalised Policy Improvement with Geometric Policy Composition. International Conference on Machine Learning (ICML), 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive comments and for finding our paper “novel” and “well-written”. We address each point below.\n\n> (W1) I don't think the term generalized occupancy model is appropriate here. \n\nWe name our method  “generalized occupancy model” because successor features capture a notion of state occupancy. Specifically, the successor measure is defined as $M^\\pi(s_0, s) := \\sum_t \\gamma^t p(s_t = s| s_0, \\pi)$, i.e. the discounted sum of probability of reaching state $s$ at each timestep. On the other hand, the state occupancy measure is defined as $\\rho^\\pi(s) := \\mathbb E_{s_0 \\sim p_0} \\sum_t \\gamma^tp(s_t=s | s_0, \\pi)$. Therefore, the successor measure is the state occupancy measure with initial state distribution set to $\\delta(s_0)$. Since successor feature is the integration of the feature function under the successor measure, it represents the expected feature under a state occupancy measure. That said, we realize the connection is not obvious and are open to changing the title to e.g. “Transferable Reinforcement Learning via Distributional Successor Features.” or another title that you feel might be more appropriate. \n\n> (W2.1, W2.3) In the introduction ... isn't the behavior policy not a particular policy? … Learning SFs on this dataset have the same policy dependence as GOMs.\n\nBy a “particular policy” we mean a policy that is trained to optimize a particular reward. Computing the successor feature allows one to quickly estimate its value under new rewards given the linear reward weights. However, standard SF does not provide a means to obtain an optimal policy for a new reward. That is to say, for a new reward function, we can evaluate how suboptimal the policy is, but we cannot make it better for this reward without retraining. What GOM does is we assume access to a dataset $D$, (which can indeed be represented by a behavior policy), and for any reward $r$, we can find the best policy for this reward within the support of the dataset. So while GOM cannot go beyond the behavior dataset, it is not restricted to evaluating the mean return of behavior policy. Instead, it can recover the best policy contained within the support of $D$ for arbitrary reward.\n\n> (W2.2) ... could you explain the difference between modeling the DSM over the mixture policy induced by a dataset when you have deterministic dynamics?\n\nWe realize that the characterization of DSM as having these limitations while GOMs don't is indeed a bit misleading. Conceptually, GOM is equivalent to DSM applied to a mixture policy induced by a dataset under the assumption of deterministic dynamics. *The key distinction between these two work is that our work shows DSM combined with deterministic dynamics can be used for transfer across rewards.* This insight is not shown conceptually or experimentally within the DSM paper. Rather, the DSM paper is motivated by robust zero-shot policy evaluation and risk-sensitive policy selection, a valuable contribution in a different light. *In the revision, we will position GOM as exploring new capabilities of a DSM-style framework, instead of addressing the limitations of DSM.*\n\n> (W3) The proposed method is closely related to 𝛾-models [2] and geometric horizon models [3]. \n\nWhile both GOMs and $\\gamma$-models / geometric horizon models learn the discounted occupancy measure under the dynamics, they differ in their transferability. Gamma models model the geometrically discounted state distribution of a policy trained to optimize a specific reward. While we can draw samples from this distribution to evaluate this policy under new rewards, we cannot improve the policy without re-optimization. On the other hand, GOMs learn a distribution over outcomes in the dataset, which can then be used to extract optimal policies for downstream tasks as long as it is covered by the data distribution.\n\n> (Q1) Ablation of feature type for USFA\n\nWe provide additional ablation studies of the feature choice for USFA in Table 2 of the attached pdf. On the antmaze-medium task, we found Fourier features to outperform transition, laplacian, and HILP features. We hypothesize that the periodic structure in Fourier features is suitable for structured rewards such as distance-to-goal in antmaze. \n\n> (Q2) Sweep over embedding dimensionality of FB\n\nWe provide additional ablations over the feature dimension of FB representation in Table 3 of the attached pdf. We found decreasing the feature dimension to provide a regularization effect, although decreasing it too much restricts the representation capacity.\n\n> (Q3) Section 6.4, why wasn't a comparison with USFA / FB performed here?\n\nWe omitted these baselines because we were not trying to compare GOMs to USFA / FB on this specific task. Rather, we were trying to demonstrate GOM’s ability to optimize arbitrary rewards, which goal-conditioned methods cannot. We add the comparisons to USFA/FB in Table 4 of the pdf. As expected, USFA / FB can distinguish the different human preferences by observing the rewards, though they perform slightly worse than GOM.\n\n> (Q4) How is the objective function (1) an off-policy update?\n\nWe say this is an off-policy update in the sense that we can use data from some dataset / behavior policy to find an optimal policy for a different task. That said, we understand your point about it being an on-policy update on the behavior policy, and we will add a clarification accordingly.\n\n> (Limitations) It's unclear how their method performs under any form of environment stochasticity.\n\nWhile our theoretical derivations are conducted with the deterministic MDP assumption, in practice the environments we evaluate on do contain stochasticity (lines 356-358), and we find our method to perform well. To address the deterministic assumption, we need to separately account for the stochasticity of the environment and the policy. We will add the limitation of deterministic dynamics assumption to the conclusion section of the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "GYQzQzcgCL", "reviewer": "Reviewer_tiLN", "summary": "The paper describes a method for using distributions of successor features for fast transfer across reinforcement learning tasks. The method combines the long-term predictions and fast transfer properties of successor features with a \"readout policy\" (conditioned on achieving a particular successor feature) and a way to sample good successor features for the readout policy to achieve (using diffusion models).", "questions": "No specific questions.", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "By modeling distributions over successor features, the paper argues that the model can avoid the policy-dependence of typical successor feature approaches. This is an important problem, and the paper does a good job of motivating the work.\n\nThe discussion of the model and its various components is clear.\n\nThe theoretical analysis is nice. While the theorems do not quite guarantee that the method will work well in practice, they are a good sign that the method should work under ideal conditions.\n\nThe experimental results are encouraging.", "weaknesses": "1. The paper's main argument in favor of the proposed method is that it overcomes the policy dependence of successor features, since the method is trained from a behavior dataset, rather than a particular policy.\n\nI find this argument unconvincing. I understand that the behavior policy can be non-stationary, but the empirical distribution of the dataset itself implicitly defines a single policy. The paper could do a better job explaining how this is providing more diversity of experience than typical policy-dependent successor feature approaches.\n\nIf I had to guess, I'd say that path dependence in the trajectory/policy interactions means that structured policies generate more interesting trajectories that are hard to sample from a simple stochastic policy. For example, walking in a specific direction for a long time, versus taking a random walk. If we sample from the implicit policy defined by the empirical data distribution, we would get less structured behavior.\n\n2. It's a little unclear to me which parts of the proposed method are new, and which are not. It's also unclear to me what the baseline successor feature methods would do in various places. For example, in the last paragraph of Section 4.2, maybe the paper could discuss what would normally be required to get S.F.s to work, so that it is clearer why GOMs represent a new class of models.\n\n3. In Def 5.2, it is unclear where $s_t$ comes from in the return. Is it sampled from $\\pi_\\beta$? Or from $\\pi$? And it's not immediately clear what $\\pi$ is being used for here. It would be helpful if the paper gave an overview of why we need $\\pi$ before introducing the definition.\n\nOne related concern here. This says that a single (s,a) pair is good if it gets lower return with low probability. But the probability that the policy gets high return would be a product over all the state-action pairs, right? Does that require a very strict $\\delta$?\n\n4. The experimental results only use 4 seeds, so it's hard to say if this will generalize. See Henderson et al. (2017, https://arxiv.org/abs/1709.06560) for a discussion on why this is so important.\n\n5. It would be helpful to see an example of the trajectories that are being stitched together.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive feedback! It is encouraging to hear that our paper is addressing “an important problem.” We reply to your questions below.\n\n > (W1) The paper could do a better job explaining how this is providing more diversity of experience than typical policy-dependent successor feature approaches.\n\nThe offline dataset indeed defines an implicit policy. Compared to a policy trained to optimize a specific reward function, the implicit policy has less structure and more coverage. Carrying over your analogy, a task-optimal policy could be walking in a particular direction for a long time, whereas the implicit dataset policy could be taking a random walk. Importantly, we can extract behaviors for walking towards many different directions from a dataset of random walks, just by taking a subset of actions at each step. In contrast, the action distribution of the task-optimal policy at each step is already narrow, so if the target task requires taking a step towards a different direction, the policy contains no information about it. This explains how modeling the distribution of successor features in the dataset enables transfer to various rewards. \n\n> (W2) It's a little unclear to me which parts of the proposed method are new, and which are not.\n\nOur work is built on a line of work in successor features and distributional RL. SF [1] first demonstrates the feasibility of estimating successor features using Bellman backup and evaluating a given policy under new rewards using the linearity argument. RaMP [2] propose random features as a viable choice of features for SF and use open-loop planning to avoid policy dependency. Our main contributions are (1) demonstrating the feasibility of using diffusion models to learn a distributional successor feature given a dataset, and (2) the ability to extra optimal policies for arbitrary rewards via a readout policy and guided diffusion planning. We will add this clarification to the related work section of the paper.\n\n> (W3) In Def 5.2, it is unclear where 𝑠𝑡 comes from in the return.\n\n$ \\mathbb P_{\\pi_\\beta}[Q^{\\pi_\\beta}(s,a)< \\sum_{t=1}^\\infty \\gamma^{t-1}r(s_t)\\mid s_1=s]$ indicates that $s_t$ is from trajectories generated by following the dataset actions $\\pi_\\beta$ starting from state $s$. We will add this clarification in the revision.\n\n> (W4) The experimental results only use 4 seeds, so it's hard to say if this will generalize. \n\nWe recognize the importance of having more seeds in RL experiments. Due to the compute bottleneck, we present results with 6 seeds for a subset of the main experiments in Table 6 of the attached pdf. While we found no significant variation in the results, we will update the results in the final revision with 6 seeds. \n\n> (W5) It would be helpful to see an example of the trajectories that are being stitched together.\n\nIn Fig 2 of the attached pdf, we visualize rollouts of the GOM policy on the roboverse PickPlace and BlockedDrawer tasks. The dataset only contains trajectories of the first phase (e.g. picking) or the second phase (e.g. placing). The GOM policy is able to generate a single trajectory that completes two phases. \n\nReferences:\n\n[1] André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, Hado van Hasselt, David Silver. Successor Features for Transfer in Reinforcement Learning. NeurIPS 2017.\n\n[2] Boyuan Chen, Chuning Zhu, Pulkit Agrawal, Kaiqing Zhang, Abhishek Gupta. RaMP: Self-Supervised Reinforcement Learning that Transfers using Random Features. NeurIPS 2023.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "ZwLabVqvef", "reviewer": "Reviewer_EMA6", "summary": "The paper presents a novel approach called Generalized Occupancy Models (GOMs), which aims to address the challenges of model-based RL and successor features in transferring across tasks with various reward functions. GOMs learn a distribution of successor features from a stationary dataset, enabling the quick selection of optimal actions for new tasks without suffering from compounding error. The paper provides a theoretical analysis of GOMs and demonstrates their efficacy in simulated robotics problems.", "questions": "- Why is Franka Kitchen using sparse rewards while this was changed for antmaze?\n- Is the guidance coefficient also sensitive to different rewards within the same environment or just to the environment itself?", "limitations": "The authors sufficiently discuss the limitations of their work.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "This work is overall well written and presented. It provides an interesting theoretical analysis of the proposed method, which is complemented by a decent empirical validation through experiments in various simulated robotics problems, demonstrating the practical applicability of GOMs. The paper also provides details regarding implementation alongside the codebase to enable reproducibility.", "weaknesses": "In general, the paper appears to oversell the generalization capabilities a bit. One of the main assumptions is the shared state space across tasks, which is often simply not the case even when assuming the same agent/embodiment and closely related tasks (different number of objects/obstacles, different goal specifications, etc.).\nThis is to the best of my understanding also reflected to some extend in the experiments, which seem to leverage the same task, merely changing the goal or sequence. In the experiments I would like to see actual changes of the underlying reward function of performing a completely different task in the same setting. This could for example include stacking cubes vs moving them away from each other, etc.\n\nAnother concern is the statistical significance of the presented results. The experiments only report the mean over four seeds with 1 standard deviation, which might limit the statistical significance of the presented results. A more robust experimental setup with a larger number of seeds and a thorough analysis of the standard deviation would provide a more convincing argument for the efficacy of GOMs.\n\nAs acknowledged by the authors, GOMs require a pre-specified feature function that linearly expresses rewards, which may not always be feasible or accurate for all types of environments or tasks. However, the paper does not provide any further insights why Fourier feature perform better than others and if this is a general phenomenon with the proposed method or if this is something that needs to be tuned/selected for different tasks. \n\n\nThe paper may overstate its generalization capabilities. A significant assumption underpinning the framework is the consistency of the state space across tasks. In practical scenarios, even with the same agent or embodiment, state spaces can vary significantly, such as in the number of objects or obstacles present or in task goal specification. This limitation is, to the best of the reviewers knowledge, somewhat reflected in the experimental design, which seems to focus on variations of the same task by altering goals or sequences rather than fundamentally changing the underlying reward functions. The experiments could be more convincing if they included tasks that are distinctly different within the same environment, such as contrasting tasks of stacking cubes versus moving them apart.\n\nThe statistical robustness of the results is another area of concern. The experiments report averages over only four seeds with 1 standard deviation, which may not sufficiently demonstrate the reliability of the findings. A more comprehensive experimental setup with a greater number of seeds would strengthen the case for the effectiveness of GOMs.\n\nWhile the authors recognize that GOMs depend on a pre-specified feature function capable of linearly expressing rewards, the paper only provides some analysis on why Fourier features were chosen and whether their performance is superior to other features in the appendix. This also raises the question whether Fourier features are generally applicable for the proposed method or if it is a parameter that needs to be fine-tuned or selected for different tasks/rewards. Addressing these points could further enhance the paper's contributions and provide clearer guidance on the applicability and limitations of GOMs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and finding our approach “novel”. We address your questions below.\n\n> (W1) In the experiments I would like to see actual changes of the underlying reward function of performing a completely different task in the same setting.\n\nWe conduct additional experiments on the Hopper environment adapted from RaMP. The environment features 4 tasks: forward, backward, stand, and jump, with complex reward functions that are not easily specified by goals. The dataset is a mixture of replay buffers of expert agents trained to achieve these tasks. We adapt the environment and dataset to our offline setting, relabeling the dataset with different task rewards for transfer. GOM compares favorably to representative baselines from each class of methods, achieving the highest average return across 4 tasks.\n\n> (W2) Another concern is the statistical significance of the presented results. \n\nWe recognize the importance of having more seeds in RL experiments. Due to the compute bottleneck, we present results with 6 seeds for a subset of the main experiments in Table 6 of the attached pdf. While we found no significant variation in the results, potentially due to the offline nature of the algorithms, we will update the results in the final revision with 6 seeds. \n\n> (W3) However, the paper does not provide any further insights why Fourier feature perform better than others and if this is a general phenomenon with the proposed method or if this is something that needs to be tuned/selected for different tasks.\n\nFor antmaze and kitchen, we found random Fourier features to outperform random features and learned  features. Recent work [1] shows that Fourier features can fit higher frequency functions, explaining its improved expressivity compared to standard random features. The reason it outperforms learned features is rather elusive, and we hypothesize that this is because the antmaze and kitchen rewards are rather structured. Hence it suffices to fit the reward with Fourier features, while learned features overfit to their pretraining objectives. For more complicated rewards in the Hopper tasks (Table 1 in pdf), we found a feature network pretrained with dynamics prediction objective to outperforms fourier feature. This applies to both our method and the USF baseline. \n\n> (Q1) Why is Franka Kitchen using sparse rewards while this was changed for antmaze?\n\nWe use a dense reward for antmaze because the task is long horizon in nature, making it difficult to propagate sparse reward signals via dynamic programming. The kitchen tasks, on the other hand, have shorter horizons. Moreover, they use a stagewise sparse reward, assigning a reward equal to the number of completed stages (less sparse than a task-completion reward). This is a standard reward function for these types of multistage tasks [2].\n\n> (Q2) Is the guidance coefficient also sensitive to different rewards within the same environment or just to the environment itself?\n\nWe found the same guidance coefficient to work well for all tasks within antmaze, kitchen, and hopper. Hence from our set of experiments the guidance coefficient is only sensitive to the environment itself but not the different rewards. \n\nReferences:\n\n[1] Ge Yang, Anurag Ajay, Pulkit Agrawal. Overcoming the Spectral Bias of Neural Value Approximation. ICLR 2022.\n\n[2] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, Sergey Levine. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. ArXiv 2020.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:
We thank the reviewers for their careful reading and constructive feedback. We appreciate the reviewers for finding our approach “promising” and “novel, our paper “well-presented” and “easy to read,” and our experiments “encouraging.” We address some common questions here and defer more detailed responses to individual comments in the threads.
- We provide additional results on a new Hopper environment (Fig. 1) adapted from RaMP [1]. The environment features 4 tasks: forward, backward, stand, and jump, each with complicated reward function not directly specifiable by goals. We modify the original reward function by removing the action penalty, making the reward only dependent on the state. The dataset contains trajectories from replay buffers of expert policies trained to achieve each task. For GOM and USF, we use learned features pretrained with a dynamics prediction objective, as we found it to perform better than random features. As shown in Table 1 in the attached pdf, GOM compares favorable to representative baselines from each class of methods (USF, RaMP, COMBO), achieving the highest average return across 4 tasks.
- In Section 2, we perform additional ablation experiments. 
    - For USF (Table 2), we ablate the choice of features and found that random Fourier features outperforms pretrained features with transition, Laplacian, or HILP [5] objectives.
    - For FB (Table 3), we ablate the feature dimensions and found decreasing the feature dimension to provide a regularization effect, although decreasing it too much restricts the representation capacity.
    - We ablate the effective planning horizon of GOM by reducing the discount factor $\gamma$ (Table 4). We found for shorter horizon tasks (antmaze-umaze), the optimality of planned trajectories improves as the planning horizon decreases. However, for long-horizon tasks (antmaze-medium), reducing the horizon too much hurts global optimality. 
- Fig. 2 of the attached pdf visualizes rollouts of the GOM policy on the roboverse PickPlace and BlockedDrawer tasks. The dataset only contains separate trajectories completing the first phase (e.g. picking) or the second phase (e.g. placing). The GOM policy is able to perform trajectory stitching and generate a single trajectory to complete the task. 
- Section 4 provides additional baseline comparisons on the preference antmaze task. USF and FB are able to distinguish different user preferences, although their performance is slightly worse than GOM. Another goal-conditioned learning baseline, GCSL [6], commits to one mode representing the shortest path, regardless of the user preference. 
- Several reviewers raise questions about the notion of “policy-independence.” We clarify that GOM models the distributional successor feature of the behavior policy, which is not optimal for a specific reward, but covers the optimal policies for various rewards. For any reward function, GOM is able to find the best policy within the support of the dataset. This is contrasted with related works such as SF [2, 3] and gamma models [4], which model the occupancy of a single task-optimal policy. This means they cannot transfer to new rewards without redoing policy optimization.

References:

[1] Boyuan Chen, Chuning Zhu, Pulkit Agrawal, Kaiqing Zhang, Abhishek Gupta. RaMP: Self-Supervised Reinforcement Learning that Transfers using Random Features. NeurIPS 2023.

[2] André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, Hado van Hasselt, David Silver. Successor Features for Transfer in Reinforcement Learning. NeurIPS 2017.

[3] Harley Wiltzer, Jesse Farebrother, Arthur Gretton, Yunhao Tang, André Barreto, Will Dabney, Marc G. Bellemare, Mark Rowland. A Distributional Analogue to the Successor Representation. ArXiv 2024.

[4] Michael Janner, Igor Mordatch, Sergey Levine. Gamma-Models: Generative Temporal Difference Learning for Infinite-Horizon Prediction. NeurIPS 2020.

[5] Seohong Park, Tobias Kreiman, Sergey Levine. Foundation Policies with Hilbert Representations. ICML 2024.

[6] Dibya Ghosh, Abhishek Gupta, Ashwin Reddy, Justin Fu, Coline Manon Devin, Benjamin Eysenbach, Sergey Levine. Learning to Reach Goals via Iterated Supervised Learning. ICLR 2021.

Meta-review: This paper proposes an approach to zero-shot reinforcement learning through learning the distribution of successor features in deterministic MDPs using offline data. It enables efficient computation of approximately optimal policies to achieve modeled outcomes. It can also generalize to new reward functions without additional environment interaction.
Empirical evaluation across numerous tasks shows strong performance and importantly  shows that the idea inspired by SFs can scale to more complex environments. 
Strengths include potential reduction in error accumulation and flexibility with novel rewards. 
Multiple reviewers agree to these evaluations and also applaud the clarity of writing. 
However there were concerns involving, limitations in generating unseen outcomes, scalability concerns for complex environments, and assumptions about reward function accessibility. The evaluation methodology also raises some fairness and comprehensiveness questions. 

There is a consensus inclination among all towards accept, with one reviewer being skeptical mainly due to limitations of baselines and experiments.  The rebuttal did a good job with clarifications and experiments. 
The meta reviewer also read the paper in detail and finds that it is an interesting idea with sufficient technical  evidence to be of interest to a broader audience in policy learning. 
The authors are advised to review the final comments, and update the manuscript accordingly.. In addition to including clarifications in the main paper, the authors are also advised to describe the possible limitations of this work in more detail.
Decision: Accept (poster)