Distributional Successor Features Enable Zero-Shot Policy Optimization

Authors: chuning  zhu
xinqi  wang
tyler  han
simon shaolei du
abhishek  gupta
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 00:33:17

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/.

Main paper: 
 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 on the single-task setting, learning behaviors that maximize a specific reward function. However, for practical deployment, RL agents must be able to generalize across different reward functions within an environment. For example, a robot deployed in a household setting should not be confined to a 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 an environment specified by a reward function. To achieve this type of generalization, we 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 is not revealed during Preprint. Under review. happen?" p(ψ|s) and "how can we achieve a particular outcome?" p(a|s, ψ). This is 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 a particular policy or reward. Moreover, once the task reward is observed, the model must provide a way to quickly evaluate and improve the policy since different tasks require different optimal policies.A natural approach to this problem is model-based reinforcement learning [61, 21, 62] , which learns an approximate dynamics model of the environment. Given a downstream reward function, taskoptimal 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 the previous step.In practice, however, autoregressive generation suffers from compounding error [31, 1, 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 short-horizon, low-dimensional problems.An alternative class of algorithms based on successor features (SFs) has emerged as a potential approach to transferable decision-making [3, 2] . Successor features represent the discounted sum of features for a given policy. Assuming a linear correspondence between features and rewards, policy evaluation under new rewards reduces to a simple linear regression problem. Notably, by directly predicting long-term outcomes, SFs avoid autoregressive rollouts and hence compounding error. However, the notion of successor features is deeply tied to the choice of a particular policy. This policy dependence hinders the recovery of optimal policies for various downstream tasks. Current approaches to circumvent policy dependence either maintain a set of policies and select the best one during inference [3] or randomly sample reward vectors and make conditional policy improvements [5, 54, 55] . Nevertheless, a turnkey solution to transfer remains a desirable goal.In this work, we propose a new class of models-Distributional 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 a 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 an action to realize it. DiSPOs are a new class of world models as they essentially capture the dynamics of the world and can be used to plan for optimal actions under arbitrary rewards, without facing the pitfalls of compounding error.Since multiple outcomes can follow from a particular state, and multiple actions can be taken to achieve a particular outcome, both the outcome distribution and the policy require expressive model classes to represent. We provide a practical instantiation of DiSPOs using diffusion models [23, 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 a suite of long-horizion simulated robotics domains and further show that DiSPOs provably converge to "best-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. 
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 multistep) dynamics models via supervised learning and use them for planning [11, 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 [1, 31] . Despite improvements to model architectures [20, 25, 1, 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) [3, 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, as they are defined as the accumulated features when acting according to a particular policy. This makes extracting optimal policies for new tasks challenging.To circumvent this policy dependence, generalized policy improvement [3, 2] 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 [6] and Forward-Backward Representations [54, 55] randomly sample reward weights z and jointly learn successor features and policies conditioned on z. 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 a 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 a concurrent work that learns a distribution over successor representations using tools from distributional RL [4] . Importantly, DSM models the distributional successor measure of a 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.e., the behavior policy), where the distribution stems from the range of meaningfully distinct long-term outcomes. This type of modeling allows DiSPOs to extract optimal behavior for arbitrary downstream tasks, while DSMs suffer from the same policy dependence that standard successor feature-based methods do. 
preliminaries: 
 We adopt the standard Markov Decision Process (MDP) notation and formalism [24] for an MDP M = (S, A, r, γ, T , ρ 0 ), but restrict our consideration to the class of deterministic MDPs. While this does not encompass every environment, it does capture a significant set of problems of practical interest. Hereafter, we refer to a 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 × A → S 1 The difference between tasks only lies in having different state-dependent Markovian reward functions r : S → [0, 1].Value Functions and Successor Features Let R = ∞ t=1 γ t-1 r(s t ) denote the cumulative reward for a trajectory {s i , a i } ∞ i=1 . One can then define the state value function under policy π as V π (s) := E π,T [R | s 1 = s], and the state-action value function asQ π (s, a) := E π,T [R | s 1 = s, a 1 = a].The value function admits a temporal structure that allows it to be estimated using dynamic programming, which iteratively applies the Bellman operator until a fixed point is reachedV π (s) := r(s) + γ E π [V π (s 2 ) | s 1 = s].While these Bellman updates are in the tabular setting, equivalent function approximator variants (e.g., with neural networks) can be instantiated to minimize a Bellman "error" with stochastic optimization techniques [37, 19, 36] .Successor features [2] generalize the notion of a value function from task-specific rewards to taskagnostic features. Given a state feature function ϕ : S → R d , the successor feature of a policy is defined as ψ π (s) = E π,T ∞ t=1 γ t-1 ϕ(s i ) | s 1 = s . Suppose rewards can be linearly expressed by the features, i.e. there exists w ∈ R n such that R(s) = w ⊤ ϕ(s), then the value function for the particular reward can be linearly expressed by the successor feature V π (s) = w ⊤ ψ π (s). Hence, given the successor feature ψ π 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 ψ π (s) := ϕ(s) + γE π [ψ π (s ′ )] , allowing them to be estimated using dynamic programming [3] . In this paper, we also refer to the discounted sum of features along a 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 a class of generative models where data generation is formulated as an iterative denoising process. Specifically, DDPM [23] consists of a forward process that iteratively adds Gaussian noise to the data, and a 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(x|y) by adding a guidance ∇ x log p(y|x) to the score function in each sampling step [13] . As we show in Sec. 4.3, 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 i , a i , s ′ i )} N i=0 of transition tuples collected with some behavior policy π β under dynamics T . The goal is to quickly obtain the optimal policy π * 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 RL [34] ), we will aim to find the best policy within dataset coverage for the downstream task. This is defined more precisely in Section 4.2. 
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 4.2. Finally, we describe a practical instantiation of DiSPOs in Section 4.3.To transfer and obtain optimal policies across different reward functions, generalist decision-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 a particular a state feature function ϕ(s), DiSPOs model the distribution of successor features p(ψ|s) over all paths that have coverage in the dataset. In deterministic MDPs, each successor feature ψ (discounted sum of features ψ = t γ t-1 ϕ(s t )) 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.e., there exists w ∈ R n such that r(s) = w ⊤ ψ(s) [42, 43, 66, 56, 9, 3] , the value of each outcome can be evaluated as w ⊤ ψ. That is, knowing w effectively transforms the distribution of outcomes p(ψ|s) into a distribution of task-specific values (sum of rewards) p(R|s). Notably, since w can be estimated by regressing rewards from features, distributional evaluation on a new reward function boils down to a simple linear regression.As in off-policy RL, the outcome distribution p(ψ|s) in DiSPOs can be learned via an approximate dynamic programming update, which is similar to a distributional Bellman update [4] :EQUATIONIntuitively, this update suggests that the distribution of successor features p θ (ψ|s) at state s maximizes likelihood over current state feature ϕ(s) added to sampled future outcomes ψ s ′ . This instantiates a fixed-point procedure, much like a distributional Bellman update. An additional benefit of the dynamic programming procedure is trajectory stitching, where combinations of subtrajectories in the dataset will be represented in the outcome distribution.(2) Readout policy: 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|s, ψ) that actualizes a desired long-term outcome ψ into the action a to be taken at state s. Along with the outcome model p θ (ψ|s), the readout policy π ρ (a|s, ψ) can be optimized via maximum-likelihood estimation: This update states that if an action a at a state s leads to a next state s ′ , then a should be taken with high likelihood for outcomes ψ, which are a combination of the current state feature ϕ(s) and future outcomesEQUATIONEQUATIONψ s ′ ∼ p θ (•|s ′ ).The outcome distribution p(ψ|s) can be understood as a natural analogue to a value function, but with two crucial differences: (1) it represents the accumulation of not just a single reward function 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 enables the selection of optimal actions for new rewards rather than being restricted to a particular (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 autoregressive generation, and that optimal actions can be obtained without additional policy optimization.To synthesize optimal policies for novel downstream reward functions using DiSPOs, two subproblems must be solved: (1) inferring the suitable linear reward weights w r for a particular reward function from a set of (s, r) tuples and (2) using the inferred w r to select an optimal action a * at a 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 r are known (i.e., r(s) = w T r ϕ(s)), the distribution of returns in the dataset p(R|s) is known through linearity. However, in most cases, rewards are not provided in functional form, making w r unknown a priori. Instead, given a dataset of D = {(s, r)} tuples, w r can be obtained by solving a simple linear regression problem arg min wr 1 |D| (s,r)∈D ∥w T r ψ(s) -r∥ 2 2 . Generating task-specific policies via distributional evaluation. Given the inferred w r and the corresponding future return distribution p(R|s) obtained through linear scaling of p(ψ|s), the optimal action can be obtained by finding the ψ with the highest possible future return that has sufficient data-support:EQUATIONwhere ϵ > 0 is a hyperparameter to ensure sufficient coverage for ψ. This suggests that the optimal outcome ψ * is the one that provides the highest future sum of rewards w T r ψ * while being valid under the environment dynamics and dataset coverage. Once a DiSPO is learned, the optimal action can be obtained by performing reward regression and searching for the optimal outcome under the dynamics to decode via the policy.This optimization problem can be solved in a number of ways. The most straightforward is via random shooting [53] , which samples a set of ψ from p(ψ|s) and chooses the one with the highest w T r ψ. Sec. 5 bases our theoretical analysis on this technique. Sec. 4.3 shows that for outcome models instantiated with diffusion models, the optimization problem can be simplified to guided diffusion sampling.Once ψ * has been obtained, the action to execute in the environment can be acquired via the readout policy π ρ (a|s, ψ * ). Fig. 3 shows the full policy optimization procedure, and we refer the reader to Appendix. F 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.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 an expressive state feature that linearly expresses a broad class of rewards. We choose the state feature ϕ 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 an outcome. To this end, we parametrize both the outcome model p(ψ|s) and the readout policy π(a|s, ψ) using a conditional diffusion model [23, 47] . We then train these models (optimize Equation 1, 2) by denoising score matching, a surrogate of maximum likelihood training [48] .Remarkably, when p(ψ|s) is parameterized by a diffusion model, the special structure of the optimization problem in Eq. 3 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 a penalized objective L = w T r ψ + α(log p(ψ|s) -log ϵ). Taking the gradient yieldsEQUATIONThe expression for ∇ ψ L(ψ, α) is simply the score function ∇ ψ log p(ψ|s) in standard diffusion training (Section 3), with the linear weights w r added as a guidance term. Planning then becomes doing stochastic gradient Langevin dynamics [57] to obtain an optimal ψ * sample, using ∇ ψ L(ψ, α) as the gradient. Guided diffusion removes the need for sampling a set of particles. As shown in Appendix E, it matches the performance of random shooting while taking significantly less inference time. In 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 a theoretical understanding of DiSPOs, we conduct an error analysis to connect the error in estimating the ground truth p 0 (ψ | 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 p 0 (ψ | s). Condition 5.1. We say the learnt outcome distribution p is an ϵ-goodapproximation if ∀ s ∈ S, ∥p(ψ | s) -p 0 (ψ | s)∥ ∞ ≤ ϵ.Since DiSPOs capture the outcome distribution of the behavior policy π β , we need a definition to evaluate a policy π with respect to π β . Definition 5.2. We say a state-action pair (s, a) is (δ, π β )-good if over the randomness of π β ,P π β [Q π β (s, a) < ∞ t=1 γ t-1 r(s t ) | s 1 = s] ≤ δ. Furthermore, if for all state s, (s, π(s)) is (δ, π β )- good, then we call π a (δ, π β )-good policy.We proceed to use Definition 5.2 to characterize the suboptimality of the DiSPO policy. Let τ denote the sampling optimality of the random shooting planner in Sec. 4.2. Specifically, we expect to sample a top τ outcome ψ from the behavior policy in O( 1 τ ) samples, whereP π β [w T r ψ ≤ t γ t-1 r(s t )] ≤ τ .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 ϵ-good outcome distribution p, the policy π given by the random shooting planner with sampling optimality τ is a (ϵ + τ, π β )-good 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 p 0 (ψ | s) will be amplified by an O 1 1-γ multiplicative factor. Corollary 5.4. If we have λ-Lipschitzness near the optimal policy, i.e., Q * (s, a * ) -Q * (s, a) ≤ λδ when (s, a) is (δ, β)-good, the suboptimality of output policy π isV * 0 (s 0 ) -V π 0 (s 0 ) ≤ λ 1-γ (τ + ϵ).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 × S| < ∞, and ∀(s, a) ∈ S × A, N (s, a, T (s, a)) ≥ 1, DiSPOs are guaranteed to 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 error 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-stitching" to actualize outcomes that combine different subtrajectories?We answer these questions through a number of experimental results in simulated robotics problems. We defer detailed descriptions of domains and baselines to Appendix D and C, as well as detailed ablative analysis to Appendix E.Antmaze [15] is a navigation domain that involves controlling a quadruped to reach some designated goal location. Each task corresponds to reaching a different goal location. We use the D4RL dataset for pretraining and dense rewards described in Appendix D for adaptation.Franka Kitchen [15] is a manipulation domain where the goal is to control a 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. ] is a 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 [9] for pretraining and shaped rewards for adaptation.Preference Antmaze is a variant of D4RL Antmaze [15] where the goal is 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 a custom dataset and design reward functions for each preference.Roboverse [46] is a tabletop manipulation environment with a robotic arm completing multi-step problems. Each task consists of two phases, and the offline dataset contains separate trajectories of each phase but not full task completion. A sparse reward is assigned to each time step of task completion.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 z and jointly learning a successor feature predictor ψ z and a policy π z conditioned on z. ψ z captures the successor feature of π z , while π z is 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 a backward representation. RaMP [9] removes the policy dependence of SF by predicting cumulative features from an initial state and an open-loop sequence of actions, which can be used for planning.Model-Based RL We compare with two variants of model-based reinforcement learning. MOPO [64] is a model-based offline RL method that learns an ensemble of dynamics models and performs actor-critic learning. COMBO [63] introduces pessimism into MOPO by training the policy using a conservative objective [30] . Goal-Conditioned RL Goal-conditioned RL enables adaptation to multiple downstream goals g. However, it is solving a more restricted class of problems than RL as goals are less expressive than rewards in the same state space. Moreover, standard GCRL is typically trained on the same set of goals as in evaluation, granting them privileged information. To account for this, we consider a goal-conditioned RL baseline GC-IQL [40, 29] and only train on goals from half the state space to show its fragility to goal distributions. We include the original method trained on test-time goals in Appendix E.6.3 Do DiSPOs enable zero-shot policy optimization across tasks? We evaluate DiSPOs on transfer problems, where the dynamics are 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 reward by subsampling a small number of transitions from the offline dataset, relabeling them with the test-time rewards, and performing linear least squares regression. While DiSPOs in principle can identify the task reward from online experience, we evaluate in the offline setting to remove the confounding factor of exploration.COMBO DiSPOsTable 1 reports 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 goalconditioned RL 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 [63] .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.While methods like goal-conditioned RL [40, 17] are restricted to shortest path goal-reaching problems, DiSPOs are able to solve problems with arbitrary reward functions. This is crucial when the reward is not easily reduced to a particular "goal". To validate this, we evaluate DiSPOs on tasks that encode nontrivial human preferences in a 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 in Fig 4 . As we see in Table 2 , DiSPOs and model-based RL obtain policies that respect human preferences and are performant for various rewards. Goalconditioned algorithms are unable to disambiguate preferences and end up with some probability of taking each path.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 exponential 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 environments 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, we choose the features as 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 a 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.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 to a 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, or as supplemental material. 
discussion: 
 This work introduced Distributional Successor Features for Zero-Shot Policy Optimization (DiSPOs), a 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 of our work open future research opportunities. First, DiSPOs require a choice of features ϕ(s) 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; a potential future direction could apply this paradigm to online adaptation, where the reward is inferred from online interactions.We provide the complete proofs of theorems and corollaries stated in Sec. 5. Throughout the following sections, we use 1 E to denote the indicator of event E.To state Theorem 5.3 rigorously, we introduce the basic setting here. Without loss of generality, let the feature at time step i, ϕ(si ) ∈ [0, 1 -γ] d and the outcome ψ = ∞ i=0 γ i ϕ(s i ) ∈ [0, 1] d .Moreover, we have a readout policy π s.t. â ∼ π(s, ψ) always leads to a successor state s ′ s.t. p( 1 γ (ψ -ϕ(s)) | s ′ ) > 0. We simplify the policy optimization phase of DiSPOs into the following form: for a given reward weight w r , in each time step, we have 1. Infer optimal outcome ψ * = A(s, p) through random shooting;2. Get corresponding action from π, â ∼ π(s, ψ * ).where the random shooting oracle A with sampling optimality τ satisfiesw T r A(s, p) ≥ min{R | 1 [w T r ψ≥R] p[ψ | s]dψ ≤ τ }.This intuitively means that there are at most probability τ ∈ [0, 1] of the behavior policy achieving higher reward. We proceed to prove Theorem 5.3.Proof. It suffices to prove that ∀s, (s, π(s)) is at least (τ + ϵ, π β )-good. For simplicity, we denoteR := min{R | 1 [w T r ψ≥R] p[ψ | s]dψ ≤ τ }. Then P π β [w T r ∞ t=1 γ t-1 ϕ(s t ) ≥ Q π β (s, π(s))] = P ψ∼p0(•|s) [w T r ψ ≥ w T r ψ * ] = 1 [w T r ψ≥w T r ψ * )] p 0 (ψ | s)dψ ≤ 1 [w T r ψ≥ R] p 0 (ψ | s)dψ = 1 [w T r ψ≥ R] p(ψ | s)dψ + 1 [w T r ψ≥ R] [p 0 (ψ | s) -p(ψ | s)] dψ ≤ τ + 1 [w T r ψ≥ R] ϵdψ ≤ τ + ϵ.Proof. Intuitively, policy π fall behind by at most λ(τ + ϵ) at each time step:∀s 1 ∈ S, V * (s 1 ) -V π (s 1 ) = Q * (s 1 , a * ) -Q π (s 1 , â) = Q * (s 1 , a * ) -Q * (s 1 , â) + Q * (s 1 , â) -Q π (s 1 , â) ≤ λ(τ + ϵ) + Q * (s 1 , â) -Q π (s 1 , â) = λ(τ + ϵ) + γE s2∼p(s1,â) [V * (s 2 ) -V π (s 2 )] ≤ λ(τ + ϵ) + γλ(τ + ϵ) + γ 2 E s3 [V * (s 3 ) -V π (s 3 )] ≤ • • • ≤ ∞ i=0 γ(τ + ϵ) = λ 1 -γ (τ + ϵ).A.3 Proof of Theorem 5.5Proof. Under the full coverage condition, the sampling optimality τ can be set to be zero for concrete actions, and the condition p(ψ|s) > ϵ becomes p(ψ|s) > 0. Then we see that planning only focuses on the supporting set of p: p(ψ | s) = 0 if and only (s, ψ) ∈ D for some time during execution, which equals p 0 (ψ | s) = 0. This indicates that supp(p) = supp(p 0 ). Therefore the conclusion follows.In this section, we derive the guided diffusion sampling from the perspective of control as inference [33] . To start, we define the trajectory-level optimality variable O as a Bernoulli variable taking the value of 1 with probability exp(R(τ )) and 0 otherwise, where R(τ ) = T t=0 γ t r(s t ) -R max . Note we subtract the max discounted return R max to make the density a valid probability distribution. Planning can be cast as an inference problem where the goal is to sample ψ * ∼ p(ψ|O). By Bayes rule, we have p(ψ|O) ∝ p(O|ψ)p(ψ)Taking the gradient of the log of both sides, we get∇ ψ log p(ψ|O) = ∇ ψ log p(O|ψ) + ∇ ψ log p(ψ) = ∇ ψ log exp(w ⊤ ψ) + ∇ ψ log p(ψ) = ∇ ψ w ⊤ ψ + ∇ ψ log p(ψ) = w + ∇ ψ log p(ψ)This implies that we can sample from p(ψ|O) by adding the regression weights w to the score at each timestep, yielding the same guided diffusion form as in Sec. 4.2.We parameterize the random Fourier features using a randomly initialized 2-layer MLP with 2048 units in each hidden layer, followed by sine and cosine activations. For a d-dimensional feature, the network's output dimension is ⌊d/2⌋ and the final feature is a 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 as a 1-D Unet with down dimensions [256, 512, 1024] . Each layer is modulated using FiLM [41] to support conditioning.We train our models on the offline dataset for 100,000 gradient steps using the AdamW optimizer [35] with batch size 2048. The learning rate for the outcome model and the policy are set to 3e -4 and adjusted according to a cosine learning rate schedule with 500 warmup steps. We train the diffusion noise prediction network with 1000 diffusion timesteps and sample using the DDIM [47] 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 a 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 (α = 0.5), Franka Kitchen (α = 0.01), and Roboverse (α = 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 the mean and standard deviation in the tables. Each experiment (pretraining + adaptation) takes 3 hours on a single Nvidia L40 GPU.Successor Features We use the implementation of Universal Successor Features Approximators [5] and Forward-Backward Representation [54, 55] from the code release of [55] . We choose random 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 1 million gradient steps and adapted using the same reward estimation method as DiSPOs.RaMP 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.Model-based RL We use the original implementations of MOPO [64] and COMBO [63] in our evaluations. We pretrain only the trainsition model on the offline transition datasets. To transfer to downstream tasks, we 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] . To remove the privileged information, we modify the sampling distribution to only sample from half of the goal space excluding the test-time goal location.Antmaze [16] is a navigation domain that involves controlling an 8-DoF quadruped robot to reach some designated goal location in a maze. Each task corresponds to reaching a different goal location. We use the standard D4RL offline dataset for pretraining. For downstream task adaptation, we replace the standard sparse reward 1(s = g) with a dense reward exp(-||s -g||2 Franka Kitchen [18] is a manipulation domain where the goal is to control a Franka arm to interact with appliances in the kitchen. Each task corresponds to interacting with a set of items in no particular order. We use the standard D4RL offline dataset for pretraining. For downstream task adaptation, we use the Markovian sparse rewards, where at each timestep the robot gets a reward equal to the number of completed tasks. We report the number of tasks completed throughout the entire episode in Table 1 .Hopper Hopper [7, 9] is a locomotion domain where the task is to control a hopper to perform various tasks, including hopping forward, hopping backward, standing, and jumping. We design a shaped reward for each task. We pretain on the offline dataset from [9] , collected by taking the replay buffers of expert SAC [19] agents trained to solve each task.Preference Antmaze is a variant of D4RL Antmaze [15] where the goal is to reach the top right cell from the bottom left cell in a custom maze shown in Fig. 4 . The two tasks in this environment are the two paths to reaching the goal, simulating different human preferences. To construct the dataset, we collect 1 million transitions using the D4RL waypoint controller. For each preference, we design a reward function that encourages the agent to take one path and not the other. Fig. 6 visualizes the dataset and the reward function for each preference.Roboverse [46] is a tabletop manipulation environment consisting of a 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 use the standard sparse reward, assigning a reward of 1 for each timestep the task is completed.To understand the impact of various design decisions on the performance of DiSPOs, we conducted systematic ablations on the various components, using the Antmaze medium-diverse as a test bed.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 it privileged information. Despite this, DiSPOs are competitive with the oracle in most domains.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 in the random shooting planner to improve planning speed, this comes at the cost of optimality. To understand the importance of feature dimension and type, we compare variants of our method that use lower-dimensional 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 featuresWe investigate the effect of dataset coverage on the performance of our method. We compare DiSPOs trained on the full D4RL dataset against two variants, one where we 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 as dataset coverage degrades.In Fig. 8 , we ablate the effective planning horizon of DiSPOs by controlling the discount factor γ. We found that for shorter horizon tasks (umaze-diverse), the optimality of planned trajectories improves as the planning horizon decreases. However, for long-horizon tasks (medium-diverse), reducing the planning horizon too much incurs cost on global optimality.To confirm the intuition of our method, we implemented a nonparametric baseline that constructs an empirical estimate of the outcome distribution. First, we calculate the empirical discounted sum of features along trajectories in the dataset. Given the set of empirical (s, a, ψ) pairs and the downstream reward weight w, we can then select the optimal action at state s through the following steps: (1) query the k nearest neighbors of s, (2) evaluate their corresponding values w ⊤ ψ, (3) take the action of the top-valued neighbor.We found this baseline to perform surprisingly well on antmaze-medium-diverse-v2. While it does 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.We visualize the stitched trajectories for roboverse environments in Fig. 7 . The dataset only contains trajectories from each phase separately, but DiSPOs can generate full trajectories by stitching the subtrajectories. Infer optimal ψ * = arg max ψ w ⊤ ψ s.t. p θ (ψ|s) > ϵ.Sample optimal action a * ∼ π(a|s, ψ * ).Execute action in the environment and add transition to buffer. 9: end for Algorithm 4 DiSPO Inference (Random Shooting) 1: Given p θ (ψ|s), π ρ (a|s, ψ), regression weight w, current state s.2: Sample N outcomes {ψ i } N i=1 ∼ p θ (ψ|s). 3: Compute corresponding values {v i } N i=1, where v i = w ⊤ ψ i . 4: Take optimal cumulant ψ * = ψ i , where i = arg max i {v i } N i=1 . 5: Sample optimal action a * ∼ π(a|s, ψ * ). Compute noise at timestep ϵ = ϵ θ (ψ t , t, s).Update noise ϵ ′ = ϵ -β √ 1 -ᾱt w.Sample next timestep action ψ t+1 using ϵ ′ . 7: end for 8: Sample optimal action a * ∼ π ρ (a|s, ψ T ).Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We include the code in supplemental materials and will release the code if the paper is accepted.• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?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, confidence intervals, or statistical significance tests, at 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 run 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. Justification: 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, it may have an impact when deployed in real-world settings where modeling is important, such as robotics or autonomous driving. In these cases, a natural challenge with GOMs is 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 to interpret the various outcomes that are modeled by GOMs.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or 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 considerations. • The conference expects that many papers will be foundational research and not 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 correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: The paper does not involve the release of models that have a high risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 a best faith effort. 12. Licenses for existing assets Question: 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? Answer: [Yes] Justification: The paper credits the owners of codebases and datasets that are used in the experiments.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• 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. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are 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 creators. • 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.For simplicity, we also use T (s, a) to denote the next state./20) to mitigate the challenge of sparse reward in long-horizon problems. 
 
 Appendix:  


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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.

3. 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.

The evaluations present might be insufficient to fully judge the quality of the approach:

4. 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.

5. 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.
questions: 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?

Is 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: 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.

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. 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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:

- 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.
- 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:
	- 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?
	- "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.
	- "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.
- 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.
- I believe the experimental methodology has some room for improvement, especially around the comparison to USFA and FB. See my questions below.

[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)

[2] Michael Janner, Igor Mordatch, and Sergey Levine. Gamma-Models: Generative Temporal Difference Learning for Infinite-Horizon Prediction. Neural Information Processing Systems (NeurIPS), 2020.

[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.
questions: 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.
- 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.
- 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?
- 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.

[1] Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does Zero-Shot Reinforcement Learning Exist? International Conference on Learning Representations (ICLR), 2023.

[2] Park Seohong, Tobias Kreiman, and Sergey Levine. Foundation Policies with Hilbert Representations. International Conference on Machine Learning (ICML), 2024.
limitations: 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.
paper_summary: paper_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).
strength_weakness: strength_weakness: 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.

I 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.

If 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.

2. 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.

3. 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.

One 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$?

4. 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.

5. It would be helpful to see an example of the trajectories that are being stitched together.
questions: questions: No specific questions.
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.).
This 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.

Another 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.

As 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. 


The 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.

The 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.

While 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.
questions: questions: - Why is Franka Kitchen using sparse rewards while this was changed for antmaze?
- Is the guidance coefficient also sensitive to different rewards within the same environment or just to the environment itself?
limitations: limitations: The authors sufficiently discuss the limitations of their work.

Comments: comment: Thank you for recognizing our contribution to the field, and for engaging with us to help make the work properly scoped and positioned! Emphasizing "zero-shot policy optimization" in the title makes a lot of sense and will help distinguish our paper from related work. As a potential renaming, we propose changing the title to "Zero-Shot Policy Optimization via Distributional Successor Features" and naming our method "Diffusion Distributional Successor Features (DDSF)." We also agree on the comparison with GHM, and will emphasize that our key contribution is the ability to do zero-shot policy optimization. We will be sure to incorporate the additional experiments and discussions to the final revision.
Thank you for running these additional seeds. I know this can be a hassle, but I appreciate that you did it anyway, and that you have promised to do the same for the remaining experiments. I will increase my score.
I thank the authors for their rebuttal. Significant effort has been put into the rebuttal and with the proposed modifications I believe this paper can be a strong contribution to a burgeoning field that studies what I'll call zero-shot policy optimization. I respond to the author(s) rebuttal below but I want to summarize the primary modification required for me to fight for this paper: remove the term generalized occupancy model. As I outline below it's confusing for multiple reasons and doesn't highlight the unique aspect of this work: the ability to perform zero-shot policy optimization in deterministic MDPs by modeling the distributional successor measure or a more computationally tractable objective in modeling distributional successor features.

I want to raise my score to an 8 based on the proposed modifications making it to the final version of the paper as well as the additional empirical results exploring different base features and a more thorough investigation of the FB baseline. I strongly believe with these modifications this paper makes for a strong contribution to the community.

---

(W1): It's not that I didn't understand the connection, it's that I believe the current naming is misleading if we are modeling features instead of states, especially in light of the distributional successor measure. I would like to see the title and naming changed for various reasons (more described below). On "Transferable Reinforcement Learning via Distributional Successor Features." I think this still undersells the work, I like the emphasis on distributional successor features but I think emphasizing that you're performing "zero-shot policy optimization" is important. When I see this title I immediately think of the title of the SF paper in which case transfer can mean policy evaluation instead of policy optimization.

(W2.1, W2.3): Thanks for this explanation, I think I'm being pedantic concerning naming at this point, when I see occupancy the first question in my mind is: occupancy of what? I'm immediately thinking about a policy. I STRONGLY suggest changing the naming of the method, not only because of my point above but to highlight what's unique here. The fact your method can do "zero-shot policy optimization" is what's important so having planning or policy in the name seems important to help clear up confusion.

(W2.2): I really like the idea of discussing your method as a new capability of the DSM, I think this will help clear up some confusion and better highlight the contributions of this work.

(W3): I feel like we took one step forward and now one step back. I don't like this framing of GOM vs. gamma models / GHMs. Again, the unique thing about this work is the ability to perform zero-shot policy optimization under deterministic dynamics and how your specific choice of a GHM leads to an efficient planning routine. The cross-entropy TD loss in the GHM paper is essentially identical to the generative modeling component in your work, the only thing that differs is the model and that you're modeling features instead of states.

(Q1, Q2, Q3): I thank the authors for these additional experiments, they helped solidify the strong performance of your method.
We understand your concern and have updated the main results for GOM and Universal SF with 10 seeds. As shown in the following table, we don’t see a significant difference in terms of the mean and the standard deviation of the runs compared to Table 1. We hope this can alleviate your concern about our results, and we will be sure to run 10 seeds for all experiments in the revision.
| Task                | GOM         | USF         |
|---------------------|-------------|-------------|
| umaze-v2            | 591 ± 15    | 458 ± 6     |
| umaze-play-v2       | 576 ± 14    | 443 ± 5     |
| medium-diverse-v2   | 648 ± 54    | 373 ± 43    |
| medium-play-v2      | 625 ± 52    | 390 ± 34    |
| large-diverse-v2    | 345 ± 50    | 233 ± 28    |
| large-play-v2       | 327 ± 67    | 237 ± 39    |
| kitchen-partial     | 40 ± 7      |   1 ± 1     |
| kitchen-mixed       | 46 ± 9      |  9 ± 9       |
We understand your concern and have updated the main results for GOM and Universal SF with 10 seeds. As shown in the following table, we don’t see a significant difference in terms of the mean and the standard deviation of the runs compared to Table 1. We hope this can alleviate your concern about our results, and we will be sure to run 10 seeds for all experiments in the revision.
| Task                | GOM         | USF         |
|---------------------|-------------|-------------|
| umaze-v2            | 591 ± 15    | 458 ± 6     |
| umaze-play-v2       | 576 ± 14    | 443 ± 5     |
| medium-diverse-v2   | 648 ± 54    | 373 ± 43    |
| medium-play-v2      | 625 ± 52    | 390 ± 34    |
| large-diverse-v2    | 345 ± 50    | 233 ± 28    |
| large-play-v2       | 327 ± 67    | 237 ± 39    |
| kitchen-partial     | 40 ± 7      |   1 ± 1     |
| kitchen-mixed       | 46 ± 9      |  9 ± 9       |
Thank you for responding to our rebuttal in detail. We address your questions below.

(W1) There is indeed a balance between staying within the behavior distribution and seeking potentially more optimal behavior by going out of distribution. This holds true for both model-based and GOM-style methods. In fact, we found the guided diffusion planner generating out-of-distribution outcomes when the guidance coefficients is too large, hence the importance of tuning the guidance coefficient. Alternatively, we can train an ensemble of GOMs to quantify the epistemic uncertainty and add a penalty term to the planning cost to generate more in-distribution behavior. We do not claim GOMs to suffer from less limitations than model-based RL. Rather, they face the same tradeoff between optimality and accuracy, which can be addressed in a similar fashion.

(W2) Thanks for clarifying your question. The overhead of random shooting @ 10 over guided diffusion comes from forwarding a 10 times larger batch through the diffusion sampling process and sorting the resulting values. The overhead is small per timestep, but accumulates over the number of diffusion and planning steps. We verify that planning with 1 random shooting sample (which is essentially sampling from the behavior policy) has roughly the same wall time as guided diffusion, at 42.5 seconds across 10 runs. 

We emphasize that the tradeoff between planners is not central to the contribution of GOMs. GOM is a versatile framework that is compatible with various planners, and the particular instantiation of GOM using diffusion models opens up the possibility of guided diffusion planner as a faster alternative to random shooting.

We agree that the long-horizon planning ability is an important advantage of GOMs. We misunderstood your original question and conducted experiments with shorter planning horizons. 

(W3) In this paper, we focus on the offline setting, investigating the transfer behavior as a result of modeling capability. We evaluate our method and baseline in a controlled setting with the same offline dataset. The number of online samples required to infer the reward function is an intriguing research question, and one that we hope to explore in a future work. To speculate, given the same online interactions, our method would infer the same reward weights as RaMP, since both use linear regression for reward inference. The difference in reward inference, therefore, would lie in the exploration behavior of the learned policies. 

(W4) We realize the potential unfairness and confusion from introducing the misspecified goal-conditioned baseline. We are open to (1) presenting the full goal-conditioned baseline in Table 1 of the main paper, or (2) removing the goal-conditioned baseline from Table 1, keeping the comparison in Table 2. We hope to hear your suggestions on which would improve the paper more.

(W5) We are glad to hear that you found the additional experiments valuable. In order to convert the Hopper task from RaMP to the GOM setting, we made two modifications: (1) the original version computes velocity by taking the finite difference of adjacent timesteps, which are not accessible from the state vector offline. We instead directly use the qvel from the state vector. (2) GOM require rewards to be a function of state, so we remove the action penalty from the reward. While this intuitively should increase the reward, the removal of action penalty can lead to behavior with higher variance. We will add these clarifications to the revision of the paper.

We hope we have addressed your questions. Let us know if you have questions regarding our response.
> (W1)

I am not sure I would fully agree with that characterization. For example, from MOPO, one of the model-based baselines the approach is compared with:

    "For the algorithm to perform reliably, it’s crucial to balance the return and risk: 1. the potential gain in performance by escaping the behavioral distribution and finding a better policy, and 2. the risk of overfitting to the errors of the dynamics at regions far away from the behavioral distribution."

Part of the evaluations in that paper is also focusing on "generalization to out-of-distribution behaviors".

> (W2, Q2)

My main focus was on how much better the proposed guided diffusion procedure is than simple random shooting. To add to my previous comment, in Table 5 in the supplementary material the performance of the proposed method and random shooting @ 100 are within each other's standard deviations (631 +/- 67 vs. 619 +/- 90). The proposed method is slightly faster (42.9s vs 58.6s), but random shooting also seems to have a significant overhead (55.5s for 10 vs 58.6s for 100) and it is unclear where it comes from.

Additionally, one of the main selling points of the approach and differentiator from model-based approaches is the ability to model long-term outcomes. If the approach is instead used in an MPC fashion with a shorter horizon, differentiation from model-based methods becomes less clear as error accumulation might be less of an issue with a shorter horizon. It would be good to do a more extensive analysis of the proposed approach and model-based methods, both used in an MPC fashion, for a range of different planning horizons.

> (W3)

The main issue is the amount of (s, r) pairs needed for the new reward. While the behavior policy could be used to collect this data in an online fashion, if the amount of data needed is comparable to the size of the original dataset that would be a significant downside of the approach. Based on results presented we cannot say how much data is actually necessary to collect. In evaluation of some other online methods (like for example RaMP) we can clearly see how performance depends on the amount of interaction data collected.

> (W4)

I agree about the comparison being delicate. One issue is that the goal-conditioned method is given an arbitrarily more difficult task (removing goal labeling from half of the state space). No other method is given that variant, so it is unclear why comparing the two numbers would make any sense. One could choose any other percentage of goal labeling to be withheld to make the method arbitrarily worse.

> (W5)

I thank the authors for pointing out that the evaluation in Table 3 also requires trajectory stitching. I would suggest making this point in the main text of the paper. Section 6.5 being explicitly about trajectory stitching and not referring to those results, but to much simpler stitching evaluation, may cause the reader to miss the point.

Additional evaluations on non-goal-reaching tasks are very much appreciated. I noticed that the performance of RaMP on "Stand" and "Jump" tasks seems lower than in the original paper (3767 ± 94 vs. ~5200 and 3098 ± 61 vs. ~4700), it is unclear why that is the case.

> (Q1)

Appreciate your response. I think that is a fair answer.

> (Limitations)

Once more, the additional evaluations are very much appreciated. Note/question about it in (W5) above.
Thank you for your detailed response and for conducting the additional experiments. While I acknowledge the improvements and the potential benefits highlighted, I fully agree with reviewer tiLN's comment on the limited sample size and therefore I will not increase my score at this time.
Thanks for the responses and for running the additional experiments. I think the benefits of accepting outweigh the risks, especially in light of the new experiments, but I won't increase my score due to the (still very low) number of seeds.

Using 6 seeds is hardly any better than 4. I strongly urge you to ensure that these results hold for at least 10 seeds, and ideally 20-30. It will only strengthen the paper. Multiple reviewers made this point, and your future readers (who you also have to convince) will have the same skepticism.

If compute is a bottleneck, that's an argument for running different experiments that are less computationally expensive---not an argument for lowering the standard of what counts as evidence. It is your responsibility to make your argument convincing, and 4-6 seeds are not convincing.