Distributional Successor Features Enable Zero-Shot Policy Optimization

Thank you for your positive feedback and finding our approach “novel”. We address your questions below.

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

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

(W2) Another concern is the statistical significance of the presented results.

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

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

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

(Q1) Why is Franka Kitchen using sparse rewards while this was changed for antmaze?

We 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].

(Q2) Is the guidance coefficient also sensitive to different rewards within the same environment or just to the environment itself?

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

References:

[1] Ge Yang, Anurag Ajay, Pulkit Agrawal. Overcoming the Spectral Bias of Neural Value Approximation. ICLR 2022.

[2] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, Sergey Levine. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. ArXiv 2020.

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.

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.

(W1) The paper could do a better job explaining how this is providing more diversity of experience than typical policy-dependent successor feature approaches.

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

(W2) It's a little unclear to me which parts of the proposed method are new, and which are not.

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

(W3) In Def 5.2, it is unclear where 𝑠𝑡 comes from in the return.

$\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.

(W4) The experimental results only use 4 seeds, so it's hard to say if this will generalize.

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

(W5) It would be helpful to see an example of the trajectories that are being stitched together.

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

References:

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

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

Thank you for your constructive comments and for finding our paper “novel” and “well-written”. We address each point below.

(W1) I don't think the term generalized occupancy model is appropriate here.

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

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

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

(W2.2) ... could you explain the difference between modeling the DSM over the mixture policy induced by a dataset when you have deterministic dynamics?

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

(W3) The proposed method is closely related to 𝛾-models [2] and geometric horizon models [3].

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

(Q1) Ablation of feature type for USFA

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

(Q2) Sweep over embedding dimensionality of FB

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

(Q3) Section 6.4, why wasn't a comparison with USFA / FB performed here?

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

(Q4) How is the objective function (1) an off-policy update?

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

(Limitations) It's unclear how their method performs under any form of environment stochasticity.

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

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.

Thank you for your valuable feedback and for finding our approach “promising.” We address your comments below.

(W1) Model-based RL can generate novel action trajectories whereas GOM is constrained to the dataset distribution.

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

(W2, Q2) Concern about planning efficiency.

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

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

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

(W4) Goal-conditioned RL unfair comparison.

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

(W5) Non-goal-conditioned and trajectory stitching experiments show capability but not better performance to comparable methods.

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

(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?

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

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

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

References:

[1] Sergey Levine, Aviral Kumar, George Tucker, Justin Fu. Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems. ArXiv 2020.