Title: Zero-Shot Reinforcement Learning from Low Quality Data
Venue: NeurIPS 2024
Accepted: yes
Keywords: reinforcement learning

offline reinforcement learning

unsupervised reinforcement learning

zero-shot reinforcement learning
Abstract: Zero-shot reinforcement learning (RL) promises to provide agents that can perform _any_ task in an environment after an offline, reward-free pre-training phase. Methods leveraging successor measures and successor features have shown strong performance in this setting, but require access to large heterogenous datasets for pre-training which cannot be expected for most real problems. Here, we explore how the performance of zero-shot RL methods degrades when trained on small homogeneous datasets, and propose fixes inspired by _conservatism_, a well-established feature of performant single-task offline RL algorithms. We evaluate our proposals across various datasets, domains and tasks, and show that conservative zero-shot RL algorithms outperform their non-conservative counterparts on low quality datasets, and perform no worse on high quality datasets. Somewhat surprisingly, our proposals also outperform baselines that get to see the task during training. Our code is available via the project page https://enjeeneer.io/projects/zero-shot-rl/.

Zero-Shot Reinforcement Learning from Low Quality Data

Abstract

Zero-shot reinforcement learning (RL) promises to provide agents that can perform any task in an environment after an offline, reward-free pre-training phase. Methods leveraging successor measures and successor features have shown strong performance in this setting, but require access to large heterogenous datasets for pre-training which cannot be expected for most real problems. Here, we explore how the performance of zero-shot RL methods degrades when trained on small homogeneous datasets, and propose fixes inspired by conservatism , a well-established feature of performant single-task offline RL algorithms. We evaluate our proposals across various datasets, domains and tasks, and show that conservative zero-shot RL algorithms outperform their non-conservative counterparts on low quality datasets, and perform no worse on high quality datasets. Somewhat surprisingly, our proposals also outperform baselines that get to see the task during training. Our code is available via the project page https://enjeeneer.io/projects/zero-shot-rl/ .

1 Introduction

Today's large pre-trained models generalise impressively to unseen vision [70] and language [9] tasks, but not to sequential decision-making problems. Zero-shot reinforcement learning (RL) attempts to correct this, asking, informally: can we pre-train an agent on a dataset of reward-free transitions such that it can perform any downstream task in an environment? Recently, methods leveraging successor features [5, 7] and successor measures [6, 82] have emerged as viable zero-shot RL candidates, returning near-optimal policies for many unseen tasks [83].

These works have assumed access to a large heterogeneous dataset of transitions for pre-training. In theory, such datasets could be curated by highly-exploratory agents during an upfront data collection phase [32, 10, 18, 62, 63, 35, 48]. However, in practice, deploying such agents in real systems can be time-consuming, costly or dangerous. To avoid these downsides, it would be convenient to skip the data collection phase and pre-train on historical datasets. Whilst these are common in the real world, they are usually produced by controllers that are not optimising for data heterogeneity [16], making them smaller and less diverse than current zero-shot RL methods expect.

Can we still perform zero-shot RL using these datasets? This is the primary question this paper seeks to answer, and one we address in four parts. First, we investigate the performance of existing methods when trained on such datasets, finding their performance suffers because of out-of-distribution state-action value overestimation, a well-observed phenomenon in single-task offline RL. Second, we develop ideas from conservatism in single-task offline RL for use in the zero-shot RL setting, introducing a straightforward regularizer of OOD values or measures that can be used by any zero-shot RL algorithm (Figure 1). Third, we conduct experiments across varied domains, tasks and datasets, showing our conservative zero-shot RL proposals outperform their non-conservative counterparts, and surpass the performance of methods that get to see the task in advance. Finally, we establish

VC-FB

collect

collect

%eval

Zeval

VC-FB

VC-FB

collec

VC-FB

a

zeval

QFB

zcollecy

aFB

Zeval

that our proposals do not hinder performance on large heterogeneous datasets, meaning adopting them presents little downside. We believe the ideas explored in this paper represent a step toward real-world deployment of zero-shot RL methods.

2 Preliminaries

Markov decision processes A reward-free Markov Decision Process (MDP) is defined by M = {S , A , T , γ } where S is the state space, A is the action space, T : S × A → ∆( S ) is the transition function, where ∆( X ) denotes the set of possible distributions over X , and γ ∈ [0 , 1) is the discount factor [77]. Given ( s 0 , a 0 ) ∈ S × A and a policy π : S → ∆( A ) , we denote Pr( ·| s 0 , a 0 , π ) and E [ ·| s 0 , a 0 , π ] the probabilities and expectations under state-action sequences ( s t , a t ) t ≥ 0 starting at ( s 0 , a 0 ) and following policy π , with s t ∼ T ( ·| s t -1 , a t -1 ) and a t ∼ π ( ·| s t ) . Given a reward function r : S → R ≥ 0 , the Q function of π for r is Q π r := ∑ t ≥ 0 γ t E [ r ( s t +1 ) | s 0 , a 0 , π ] .

Zero-shot RL For pre-training, the agent has access to a static offline dataset of reward-free transitions D = { ( s i , a i , s i +1 ) } |D| i =1 generated by an unknown behaviour policy, and cannot interact with the environment. At test time, a reward function r eval specifying a task is revealed and the agent must return a policy for the task without any further planning or learning. Ideally, the policy should maximise the expected discounted return on the task E [ ∑ t ≥ 0 γ t r eval ( s t +1 ) | s 0 , a 0 , π ] . The reward function is specified either via a small dataset of reward-labelled states D labelled = { ( s i , r eval ( s i )) } k i =1 with k ≤ 10 , 000 or as an explicit function s ↦→ r eval ( s ) (like 1 at a goal state and 0 elsewhere). Intuitively, the zero-shot RL problem asks: is it possible to train an agent using a pre-collected dataset of transitions from an environment such that, at test time, it can return the optimal policy for any task in that environment without any further planning or learning?

State-of-the-art zero-shot RL methods leverage either successor measures [6] or successor features [5], with the former instantiated by forward backward representations [82] and the latter by universal successor features [7]. The remainder of this section introduces these ideas.

Successor measures The successor measure M π ( s 0 , a 0 , · ) over S is the cumulative discounted time spent in each future state s t +1 after starting in state s 0 , taking action a 0 , and following policy π thereafter:

The Q function of policy π for task r is the integral of r with respect to M

π :

The forward-backward framework FB representations [82] approximate the successor measures of near-optimal policies for any task. Let ρ be an arbitrary state distribution, and R d be a representation

space. FB representations are composed of a forward model F : S × A × R d → R d , a backward model B : S → R d , and set of polices ( π z ) z ∈ R d . They are trained such that

and

Intuitively, Equation 3 says that the approximated successor measure under π z from ( s 0 , a 0 ) to s is high if their respective forward and backward embeddings are similar i.e. have large dot product. By comparing Equation 2 and Equation 3, we see that an FB representation can be used to approximate the Q function of π z with respect to any reward function r as:

Training of F and B is done with TD learning [71, 78] using transition data sampled from D :

where s + is sampled independently of ( s t , a t , s t +1 ) , ¯ F and ¯ B are lagging target networks, and Z is a task sampling distribution. The policy is trained in an actor-critic formulation [47]. See [82] for a full derivation of the TD update, and our Appendix B.1 for practical implementation details including the specific choice of task sampling distribution Z .

By relating Equations 4 and 5, we find z = E s ∼ ρ [ r ( s ) B ( s )] for some reward function r . At test time, we can use this property to perform zero-shot RL. Using D labelled, we estimate the task as z eval ≈ E s ∼D labelled [ r eval ( s ) B ( s )] and pass it as an argument to π z . If z eval lies within the task sampling distribution Z used during pre-training, then π z ( s ) ≈ arg max a Q π z r eval ( s, a ) , and hence this policy is approximately optimal for r eval .

(Universal) successor features Successor features assume access to a basic feature map φ : S ↦→ R d that embeds states into a representation space, and are defined as the expected discounted sum of future features ψ π ( s 0 , a 0 ) := E [ ∑ t ≥ 0 γ t φ ( s t +1 ) | s 0 , a 0 , π ] [5]. They are made universal by conditioning their predictions on a family of policies π z

with

Like FB, USFs are trained using TD learning on

where ¯ ψ is a lagging target network, and Z is the same z sampling distribution used for FB. We refer the reader to [7] for a derivation of the TD update and full learning procedure. Test time policy inference is performed similarly to FB. Using D labelled, the task is inferred by performing a linear regression of r eval onto the features: z eval := arg min z E s ∼D labelled [( r eval ( s ) -φ ( s ) ⊤ z ) 2 ] before it is passed as an argument to the policy.

3 Zero-Shot RL from Low Quality Data

In this section we introduce methods for improving the performance of zero-shot RL methods on low quality datasets. In Section 3.1, we explore the failure mode of existing methods on such datasets. Then, in Section 3.2, we propose straightforward amendments to these methods that address the failure mode. Finally, in Section 3.3, we illustrate the usefulness of our proposals with a controlled example. We develop our methods within the FB framework because of its superior empirical performance [83], but our proposals are also compatible with USF. We push their derivation to Appendix D for brevity.

RND

Random

30

20

10

0

1000

M

n=1000

O1)

n=10000

Return

500

n=100000

0

0

50

100

0

50

100

Learning steps (k)

Learning steps (k)

3.1 Failure Mode of Existing Methods

To investigate the failure mode of existing methods we examine the FB loss (Equation 6) more closely. The TD target includes an action produced by the current policy a t +1 ∼ π z ( s t +1 ) . Equation 4 shows this is the policy's current best estimate of the Q -maximising action in state s for task z . For finite datasets, this maximisation does not constrain the policy to actions observed in the dataset, and so it can become biased towards out-of-distribution (OOD) actions thought to be of high value. In such instances F and B are updated towards targets for which the dataset provides no support. This distribution shift is a well-observed phenomenon in single-task offline RL [42, 46, 44], and is exacerbated by small, low-diversity datasets as we explore in Figure 2.

3.2 Mitigating the Distribution Shift

In the single-task setting, the distribution shift is addressed by applying constraints to either the policy, value function or model (see Section 6 for a summary of past work). Here we re-purpose single-task value function and model regularisation for use in the zero-shot RL setting. To avoid further complicating zero-shot RL methods, we only consider regularisation techniques that do not introduce new parametric functions. We discuss the implications of this decision in Section 5.

Conservative Q -learning (CQL) [42, 44] regularises the Q function by querying OOD state-action pairs and suppressing their value. This is achieved by adding new term to the usual Q loss function

where α is a scaling parameter, µ ( a | s ) is a policy distribution selected to find the maximum value of the current Q function iterate, H ( µ ) is the entropy of µ used for regularisation, and L Q is the normal TD loss on Q . Equation 10 has the dual effect of minimising the peaks in Q under µ whilst maximising Q for state-action pairs in the dataset.

We can replicate a similar form of regularisation in the FB framework, substituting F ( s, a, z ) ⊤ z for Q in Equation 10 and adding the normal FB loss (Equation 6)

The key difference between Equations 10 and 11 is that the former suppresses the value of OOD actions for one task, whereas the latter does so for all task vectors drawn from Z . We call models learnt with this loss value-conservative forward-backward representations (VC-FB).

Because FB derives Q functions from successor measures (Equation 5), and because (by assumption) rewards are non-negative, suppressing the predicted measures for OOD actions provides an alternative route to suppressing their Q values. As we did with VC-FB, we can substitute FB's successor measure approximation F ( s, a, z ) ⊤ B ( s + ) into Equation 10, which yields:

(a)

(b)

(c)

+

?

?

+

Equation 12 has the effect of suppressing the expected visitation count to goal state s + when taking an OOD action for all task vectors drawn from Z , which says, informally, if we don't know where OOD actions take us in the MDP, we assume they have low probability of taking us to any future states for all tasks. This is analogous to works that regularise model predictions in the single-task offline RL setting [37, 96, 69]. As such, we call this variant a measure-conservative forward-backward representation (MC-FB). Since it is not obvious a priori whether the VC-FB or MC-FB form of conservatism would be more effective in practice, we evaluate both in Section 4.

Implementing these proposals requires two new model components: 1) a conservative penalty scaling factor α and 2) a way of obtaining policy distribution µ ( a | s ) that maximises the current value or measure iterate. For 1), we observe fixed values of α leading to fragile performance, so dynamically tune it at each learning-see Appendix B.1.4. For 2), the choice of maximum entropy regularisation following [44]'s CQL( H ) allows µ to be approximated conveniently with a log-sum exponential across Q values derived from the current policy distribution and a uniform distribution. That this is true is not obvious, so we refer the reader to the detail and derivations in Section 3.2, Appendix A, and Appendix E of [44], as well as our adjustments to [44]'s theory in Appendix B.1.3. Code snippets demonstrating the required changes to a vanilla FB implementation are provided in Appendix G. We emphasise these additions represent only a small increase in the number of lines required to implement existing methods.

3.3 A Didactic Example

To understand situations in which a conservative zero-shot RL methods may be useful, we introduce a modified version of Maze from the ExORL benchmark [95]. Episodes begin with a point-mass initialised in the upper left of the maze ( ⊚ ), and the agent is tasked with selecting x and y tilt directions such that the mass is moved towards one of two goal locations ( ⊛ and ⊛ ). The action space is two-dimensional and bounded in [ -1 , 1] . We take the RND dataset and remove all 'left" actions such that a x ∈ [0 , 1] and a y ∈ [ -1 , 1] , creating a dataset that has the necessary information for solving the tasks, but is inexhaustive (Figure 3 (a)). We train FB and VC-FB on this dataset and plot the highest-reward trajectories-Figure 3 (b) and (c). FB overestimates the value of OOD actions and cannot complete either task. Conversely, VC-FB synthesises the requisite information from the dataset and completes both tasks.

4 Experiments

In this section we perform an empirical study to evaluate our proposals. We seek answers to four questions: (Q1) Can our proposals from Section 3 improve FB performance on small and/or lowquality exploratory datasets? (Q2) How does the performance of VC-FB and MC-FB vary with respect to task type and dataset diversity? (Q3) Do we sacrifice performance on full datasets for

SF-LAP

GC-IQL

FB

MC-FB

VC-FB

CQL

140

1.0

Score

一

120

Task

0.8

Performance

100

IV

0.6

80

Score

Zero-shot

60

0.4

Probability

40

一

0.2

20

0

0.0

0

200

400

600

800

1000

TaskScore

performance on small and/or low-quality datasets? (Q4) If our pre-training dataset only covers behaviour related to one downstream task (i.e. the dataset distribution is narrow and not exploratory), can our proposals from Section 3 improve FB performance on that task?

4.1 Setup

Domains We respond to Q1-Q3 using the ExORL benchmark [95]. ExORL provides datasets collected by unsupervised exploratory algorithms on the DeepMind Control Suite [80]. We select three of the same domains as [83]: Walker, Quadruped and Maze, but substitute Jaco for Cheetah. This provides two locomotion domains and two goal-reaching domains. Within each domain, we evaluate on all tasks provided by the DeepMind Control Suite for a total of 17 tasks across four domains. Full details are provided in Appendix A.1. We respond to Q4 using the D4RL benchmark [21]. We select the two MuJoCo [81] environments from the Open AI gym [8] that closest resemble those from ExORL: Walker2D and HalfCheetah.

Datasets For Q1-Q3 we pre-train on three datasets of varying quality from ExORL. There is no unambiguous metric for quantifying dataset quality, so we use the reported performance of offline TD3 on Maze for each dataset as a proxy 1 . We choose datasets collected via Random Network Distillation (RND) [10], Diversity is All You Need (DIAYN) [18], and RANDOM policies, where agents trained on RND are the most performant, on DIAYN are median performers, and on RANDOM are the least performant. As well as selecting for quality, we also select for size by uniformly subsampling 100,000 transitions from each dataset. For Q4 we choose the 'medium", 'medium-replay", and 'medium-expert" datasets from D4RL, each providing different fractions of random, medium and expert task-directed trajectories. More details on the datasets are provided in Appendix A.2.

4.2 Baselines

We compare our proposals to baselines from three categories: 1) zero-shot RL methods, 2) goalconditioned RL (GCRL) methods, and 3) single-task offline RL methods. From category 1), we use the state-of-the-art successor measure based method, FB, and the state-of-the-art successor feature based method, SF with features from Laplacian eigenfunctions (SF-LAP) [83]. From category 2), we use goal-conditioned IQL (GC-IQL) [60], a state-of-the-art GCRL method that, like our proposals, regularises the value function at OOD state-actions. We condition GC-IQL on the goal state on Maze and Jaco, and on the state in D labelled with highest reward on Walker and Quadruped in lieu of a

Zero-shot Performance (abs.)

Random

100k

RND

100k

DIAYN

100k

Walker

Maze

Quadruped

Jaco

FB

MC-FB

VC-FB

CQL

SF-LAP

GC-IQL

0

20

40

0

200

400

0

200

0

200

400

0

20

40

0

200

400

0

200

0

200

400

0

20

40

0

200

400

0

200

0

200

400

well-defined goal state. From category 3), we use CQL and offline TD3 trained on the same datasets relabelled with task rewards. CQL approximates what an algorithm with similar mechanistics can achieve when optimising for one task in a domain rather than all tasks. Offline TD3 exhibits the best aggregate single-task performance on the ExORL benchmark, so it should be indicative of the maximum performance we could expect to extract from a dataset. Full implementation details for all algorithms are provided in Appendix B. The full evaluation protocol is described in Appendix A.5. Appendix A.6 provides a breakdown of the computational resources used in this work.

4.3 Results

Q1 We report the aggregate performance of our baselines and proposals on ExORL in Figure 4. Both MC-FB and VC-FB outperform the zero-shot RL and GCRL baselines, achieving 150% and 137% of FB's IQM performance respectively. The performance gap between FB and SF-LAP is consistent with the results in [83]. MC-FB and VC-FB outperform our single-task baseline in expectation, reaching 111% and 120% of CQL's IQM performance

VC-FB

MC-FB

FB

400

Return

300

IQM

Aggregate |

200

100

0

105

106

107

DatasetSize

respectively despite not having access to taskspecific reward labels and needing to fit policies for all tasks . This is a surprising result, and to the best of our knowledge, the first time a multitask offline agent has been shown to outperform a single-task analogue. CQL outperforms offline TD3 in aggregate, so we drop offline TD3 from the core analysis, but report its full results in Appendix C alongside all other methods. We note FB achieves 80% of single-task offline TD3, which roughly aligns with the 85% performance on the full datasets reported by [83].

Q2 We decompose the methods' performance with respect to domain and dataset diversity in Figure 5. The largest gap in performance between the conservative FB variants and FB is on RND. VC-FB and MC-FB reach 2.5 × and 1.8 × of FB performance respectively, and outperform CQL on three of the four domains. On DIAYN, the conservative variants outperform all methods and reach 1.3 × CQL's score. On the RANDOM dataset, all methods perform similarly poorly, except for CQL on Jaco, which outperforms all methods. However, in general, these results suggest the RANDOM dataset is not informative enough to extract valuable policies-discussed further in response to Q3. There appears to be little correlation between the type of domain (Appendix A.1) and the score achieved by any method. GC-IQL performs particularly well on the goal-reaching domains as expected, but worse than all zero-shot methods on the locomotion tasks, irrespective of whether they are conservative or not. This is presumably because the goal-state used to condition the policy ( i.e. the state with highest reward in D labelled ) is a poor proxy for the true, dense reward function.

Q3 We report the aggregated performance of all FB methods across domains when trained on the full datasets in Table 1 (a full breakdown of results in provided in Appendix C). Both conservative FB variants slightly exceed the performance of vanilla FB in expectation. The largest relative performance improvement is on the RANDOM dataset-MC-FB performance is 20% higher than FB, compared to 5% higher on DIAYN and 2% higher on RND. This corroborates the hypothesis that RANDOM-100K was not informative enough to extract valuable policies.

Table 1 and Figure 4 suggest the performance gap between the conservative FB variants and vanilla FB changes as dataset size is varied. We further explore this effect in Figure 6 where we scale the RND dataset size from 10 5 through 10 7 and plot aggregate IQM performance of FB, VCFB and MC-FB across all domains. We find that the performance gap decreases as dataset size increases. This result is to be expected: a larger dataset size for a fixed exploration algorithm means a t +1 ∼ π z ( s t +1 ) in the FB TD update (Equation 6) is more likely to be in the dataset, the policy is less likely to become biased toward OOD actions, and conservatism is less needed.

Q4 We report the aggregate performance of all zero-shot RL methods and CQL on our D4RL domains in Figure 7. FB fails all domain-dataset tasks, and reaches only 10% of CQL's aggregate

SF-LAP

MC-FB

CQL

FB

VC-FB

125

Performance

100

75

Zero-shot

50

25

0

performance. MC-FB and VC-FB improve on FB's considerably (by 5.6 × and 6.8 × respectively) but under-perform CQL. SF-LAP outperforms FB, but under-performs VC-FB, MC-FB and CQL.

5 Discussion and Limitations

Performance discrepancy between conservative variants Why does VC-FB outperform MC-FB on both ExORL and D4RL? To understand, we inspect the regularising effect of both models more closely. VC-FB regularises OOD actions on F ( s, a, z ) ⊤ z , with s ∼ D , and z ∼ Z , whilst MC-FB regularises OOD actions on F ( s, a, z ) ⊤ B ( s + ) , with ( s, s + ) ∼ D and z ∼ Z . Note the trailing z in VC-FB is replaced with B ( s + ) in MC-FB which ties its updates to D further. We hypothesised that as |D| reduces, B ( s + ) provides poorer task coverage than z ∼ Z , hence the comparable performance on full datasets and divergent performance on 100k datasets.

To test this, we evaluate a third conservative variant called directed ( D )VC-FB which replaces all z ∼ Z in VC-FB with B ( s + ) such that OOD actions are regularised on F ( s, a, B ( s + )) ⊤ B ( s + ) with ( s, s + ) ∼ D . This ties conservative updates entirely to D , and according to our above hypothesis, D VC-FB should perform worse than VC-FB and MC-FB on the 100k ExORL datasets. See Appendix B.1.6 for implementation details. We evaluate this variant on all 100k ExORL datasets, domains and tasks and compare with FB, VC-FB and MC-FB in Table 2. See Appendix C for a full breakdown.

We find the aggregate relative performance of each method is as expected i.e. D VC-FB < MC-FB < VC-FB. As a consequence we conclude that VC-FB should be preferred for small datasets with

no prior knowledge of the dataset or test tasks. Of course, for a specific domain-dataset pair, B ( s + ) with s + ∼ D may happen to cover the tasks well, and MC-FB may outperform VC-FB. We suspect this was the case for all datasets on the Jaco domain for example. Establishing whether this will be true a priori requires either relaxing the restrictions imposed by the zero-shot RL setting, or better understanding of the distribution of tasks in z -space and their relationship to pre-training datasets. The latter is important future work.

Avoiding new parametric functions State-of-the-art zero-shot RL methods are complex, and we wanted to avoid further complicating them with new parametric functions. This limited our solutionspace to CQL-style regularisation techniques, but had we relaxed this constraint, other options become available. Methods like AWAC [58], IQL [40], and X -QL [25] all require an estimate of the state-value function which is not immediately accessible in the FB or USF frameworks. In theory, we could learn an action-independent USF of the form V ( s, z ) = E [ ∑ t ≥ 0 γ t φ ( s t +1 ) | s 0 , π z ] ∀ s 0 ∈ S , z ∈ R d concurrently to F and B (or ψ for USFs). If learnt with expectile regression, this function could be used to implement IQL and X -QL style regularisation; without expectile regression it could be used to compute the advantage weighting required for AWAC. It's possible that implementing these methods could improve downstream performance and reduce computational overhead at the cost of increased training complexity. We leave this worthwhile investigation for future work. We provide detail of negative results related to downstream finetuning of FB models in Appendix E to help inform future research.

D4RL Performance Unlike the ExORL results, VC-FB and MC-FB do not outperform CQL on the D4RL benchmark. We believe these narrower data distributions require a more careful selection of the conservative penalty scaling factor α . We explore this further in Appendix F, and note this is corroborated by findings in the original CQL paper [44]. Methods described above, like IQL, have been shown to be more robust than CQL partly because they bypass α tuning. We expect that exploring the integration of these methods may improve D4RL performance.

6 Related Work

Zero-shot RL Zero-shot RL methods build upon successor representations [15], universal value function approximators [74], successor features [5] and successor measures [6]. The state-of-the-art methods instantiate these ideas as either universal successor features (USFs) [7] or forward-backward (FB) representations [82, 83], with recent work showing the latter can be used to perform a range of imitation learning techniques efficiently [65]. A representation learning method is required to learn the features for USFs, with past works using inverse curiosity modules [62], diversity methods [49, 29], Laplacian eigenfunctions [87], or contrastive learning [13]. No works have yet explored the issues arising when training these methods on low quality offline datasets, and only one has investigated applying these ideas to real-world problems [34].

Goal-conditioned RL methods train policies to reach any goal state from any other state, and so can be used to perform zero-shot RL in goal-reaching environments [60, 54, 93, 19, 85]. However, they have no principled mechanism for conditioning policies on 'dense' reward functions (as such tasks are not solved by simply reaching a particular state), and so are not full zero-shot RL methods. A concurrent line of work trains policies using sequence models conditioned on reward-labelled histories [12, 33, 45, 68, 99, 11, 24, 76, 91, 90], but, unlike zero-shot RL methods, these works do not have a robust mechanism for generalising to different reward functions as test-time.

Offline RL Offline RL algorithms require regularisation of policies, value functions, models, or a combination to manage the offline-to-online distribution shift [46]. Past works regularise policies with explicit constraints [88, 20, 23, 22, 27, 64, 43, 86, 94], via important sampling [66, 79, 50, 57, 26], by leveraging uncertainty in predictions [89, 98, 4, 36], or by minimising OOD action queries [84, 14, 40], a form of imitation learning [72, 73]. Other works constrain value function approximation

so OOD action values are not overestimated [44, 42, 52, 53, 51, 92]. Offline model-based RL methods use the model to identify OOD states and penalise predicted rollouts passing through them [97, 37, 96, 2, 55, 67, 69]. All of these works have focused on regularising a finite number of policies; in contrast we extend this line of work to the zero-shot RL setting which is concerned with learning an infinite family of policies.

7 Conclusion

In this paper, we explored training agents to perform zero-shot reinforcement learning (RL) from low quality data. We established that the existing methods suffer in this regime because they overestimate the value of out-of-distribution state-action values, a well-observed pheneomena in single-task offline RL. As a resolution, we proposed a family of conservative zero-shot RL algorithms that regularise value functions or dynamics predictions on out-of-distribution state-action pairs. In experiments across various domains, tasks and datasets, we showed our proposals outperform their non-conservative counterparts in aggregate and sometimes surpass our task-specific baseline despite lacking access to reward labels a priori . In addition to improving performance when trained on sub-optimal datasets, we showed that performance on large, diverse datasets does not suffer as a consequence of our design decisions. Our proposals represent a step towards the use of zero-shot RL methods in the real world.

References

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems , 34:29304-29320, 2021.

Arthur Argenson and Gabriel Dulac-Arnold. Model-based offline planning. arXiv preprint arXiv:2008.05556 , 2020.

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhihong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. arXiv preprint arXiv:2202.11566 , 2022.

André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. Advances in neural information processing systems , 30, 2017.

Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07123 , 2021.

Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado Van Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. arXiv preprint arXiv:1812.07626 , 2018.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 , 2016.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810.12894 , 2018.

Yevgen Chebotar, Quan Vuong, Alex Irpan, Karol Hausman, Fei Xia, Yao Lu, Aviral Kumar, Tianhe Yu, Alexander Herzog, Karl Pertsch, et al. Q-transformer: Scalable offline reinforcement learning via autoregressive q-functions. arXiv preprint arXiv:2309.10150 , 2023.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning , pages 1597-1607. PMLR, 2020.

Xinyue Chen, Zijian Zhou, Zheng Wang, Che Wang, Yanqiu Wu, and Keith Ross. Bail: Bestaction imitation learning for batch deep reinforcement learning. Advances in Neural Information Processing Systems , 33:18353-18363, 2020.

Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural computation , 5(4):613-624, 1993.

Gabriel Dulac-Arnold, Daniel Mankowitz, and Todd Hester. Challenges of real-world reinforcement learning. arXiv preprint arXiv:1904.12901 , 2019.

Bradley Efron. Bootstrap methods: another look at the jackknife. In Breakthroughs in statistics: Methodology and distribution , pages 569-593. Springer, 1992.

Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint arXiv:1802.06070 , 2018.

Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive learning as goal-conditioned reinforcement learning. Advances in Neural Information Processing Systems , 35:35603-35620, 2022.

Rasool Fakoor, Jonas W Mueller, Kavosh Asadi, Pratik Chaudhari, and Alexander J Smola. Continuous doubly constrained batch reinforcement learning. Advances in Neural Information Processing Systems , 34:11260-11273, 2021.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Scott Fujimoto, Wei-Di Chang, Edward Smith, Shixiang Shane Gu, Doina Precup, and David Meger. For sale: State-action representation learning for deep reinforcement learning. Advances in Neural Information Processing Systems , 36, 2024.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning , pages 2052-2062. PMLR, 2019.

Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:2111.10364 , 2021.

Divyansh Garg, Joey Hejna, Matthieu Geist, and Stefano Ermon. Extreme q-learning: Maxent rl without entropy. International Conference on Learning Representations , 2023.

Carles Gelada and Marc G Bellemare. Off-policy deep reinforcement learning by bootstrapping the covariate shift. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 33, pages 3647-3655, 2019.

Seyed Kamyar Seyed Ghasemipour, Dale Schuurmans, and Shixiang Shane Gu. Emaq: Expected-max q-learning operator for simple yet effective offline and online rl. In International Conference on Machine Learning , pages 3682-3691. PMLR, 2021.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning , pages 1861-1870, 2018.

Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. arXiv preprint arXiv:1906.05030 , 2019.

Charles R Harris, K Jarrod Millman, Stéfan J Van Der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J Smith, et al. Array programming with numpy. Nature , 585(7825):357-362, 2020.

John D Hunter. Matplotlib: A 2d graphics environment. Computing in science & engineering , 9(03):90-95, 2007.

Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05397 , 2016.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:12731286, 2021.

Scott Jeen, Alessandro Abate, and Jonathan M Cullen. Low emission building control with zeroshot reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 37, pages 14259-14267, 2023.

Chi Jin, Akshay Krishnamurthy, Max Simchowitz, and Tiancheng Yu. Reward-free exploration for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning , pages 4870-4879, 13-18 Jul 2020.

Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In International Conference on Machine Learning , pages 5084-5096. PMLR, 2021.

Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems , 33:21810-21823, 2020.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Yehuda Koren. On spectral graph drawing. In International Computing and Combinatorics Conference , pages 496-508. Springer, 2003.

Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. In International Conference on Machine Learning , pages 5774-5783. PMLR, 2021.

Aviral Kumar, Rishabh Agarwal, Xinyang Geng, George Tucker, and Sergey Levine. Offline qlearning on diverse multi-task data both scales and generalizes. arXiv preprint arXiv:2211.15144 , 2022.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing offpolicy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint arXiv:2006.04779 , 2020.

Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. Advances in neural information processing systems , 35, 2022.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In ICLR (Poster) , 2016.

Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. In International Conference on Machine Learning , pages 6736-6747. PMLR, 2021.

Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. In International Conference on Machine Learning , pages 6736-6747. PMLR, 2021.

Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Off-policy policy gradient with state distribution correction. arXiv preprint arXiv:1904.08473 , 2019.

Jiafei Lyu, Xiaoteng Ma, Xiu Li, and Zongqing Lu. Mildly conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 35:1711-1724, 2022.

Xiaoteng Ma, Yiqin Yang, Hao Hu, Qihan Liu, Jun Yang, Chongjie Zhang, Qianchuan Zhao, and Bin Liang. Offline reinforcement learning with value-based episodic memory. arXiv preprint arXiv:2110.09796 , 2021.

Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional reinforcement learning. Advances in Neural Information Processing Systems , 34:19235-19247, 2021.

Yecheng Jason Ma, Jason Yan, Dinesh Jayaraman, and Osbert Bastani. How far i'll go: Offline goal-conditioned reinforcement learning via f -advantage regression. arXiv preprint arXiv:2206.03023 , 2022.

Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Gu. Deployment-efficient reinforcement learning via model-based offline optimization. arXiv preprint arXiv:2006.03647 , 2020.

Wes McKinney et al. pandas: a foundational python library for data analysis and statistics. Python for high performance and scientific computing , 14(9):1-9, 2011.

Ofir Nachum, Yinlam Chow, Bo Dai, and Lihong Li. Dualdice: Behavior-agnostic estimation of discounted stationary distribution corrections. Advances in neural information processing systems , 32, 2019.

Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359 , 2020.

Mitsuhiko Nakamoto, Yuexiang Zhai, Anikait Singh, Max Sobol Mark, Yi Ma, Chelsea Finn, Aviral Kumar, and Sergey Levine. Cal-ql: Calibrated offline rl pre-training for efficient online fine-tuning. arXiv preprint arXiv:2303.05479 , 2023.

Seohong Park, Dibya Ghosh, Benjamin Eysenbach, and Sergey Levine. Hiql: Offline goalconditioned rl with latent states as actions. Advances in Neural Information Processing Systems , 37, 2023.

Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.

Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning , pages 27782787. PMLR, 2017.

Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta. Self-supervised exploration via disagreement. In International conference on machine learning , pages 5062-5071. PMLR, 2019.

Zhiyong Peng, Changlin Han, Yadong Liu, and Zongtan Zhou. Weighted policy constraints for offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 37, pages 9435-9443, 2023.

Matteo Pirotta, Andrea Tirinzoni, Ahmed Touati, Alessandro Lazaric, and Yann Ollivier. Fast imitation via behavior foundation models. In International Conference on Learning Representations , 2024.

Doina Precup, Richard S Sutton, and Sanjoy Dasgupta. Off-policy temporal-difference learning with function approximation. In ICML , pages 417-424, 2001.

Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. In Learning for Dynamics and Control , pages 1154-1168. PMLR, 2021.

Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06175 , 2022.

Marc Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 16082-16097. Curran Associates, Inc., 2022.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 10684-10695, 2022.

Arthur L Samuel. Some studies in machine learning using the game of checkers. IBM Journal of research and development , 3(3):210-229, 1959.

Stefan Schaal. Learning from demonstration. Advances in neural information processing systems , 9, 1996.

Stefan Schaal. Is imitation learning the route to humanoid robots? Trends in cognitive sciences , 3(6):233-242, 1999.

Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In International conference on machine learning , pages 1312-1320. PMLR, 2015.

Kajetan Schweighofer, Andreas Radler, Marius-Constantin Dinu, Markus Hofmarcher, Vihang Patil, Angela Bitto-Nemling, Hamid Eghbal-zadeh, and Sepp Hochreiter. A dataset perspective on offline reinforcement learning. arXiv preprint arXiv:2111.04714 , 2021.

Max Siebenborn, Boris Belousov, Junning Huang, and Jan Peters. How crucial is transformer in decision transformer? arXiv preprint arXiv:2211.14655 , 2022.

Richard Sutton and Andrew Barto. Reinforcement Learning: An Introduction . The MIT Press, second edition, 2018.

Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning , 3:9-44, 1988.

Richard S Sutton, A Rupam Mahmood, and Martha White. An emphatic approach to the problem of off-policy temporal-difference learning. The Journal of Machine Learning Research , 17(1):2603-2631, 2016.

Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00690 , 2018.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems , pages 5026-5033. IEEE, 2012.

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

Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? In The Eleventh International Conference on Learning Representations , 2023.

Qing Wang, Jiechao Xiong, Lei Han, Han Liu, Tong Zhang, et al. Exponentially weighted imitation learning for batched historical data. Advances in Neural Information Processing Systems , 31, 2018.

Tongzhou Wang, Antonio Torralba, Phillip Isola, and Amy Zhang. Optimal goal-reaching reinforcement learning via quasimetric learning. In International Conference on Machine Learning , pages 36411-36430. PMLR, 2023.

Jialong Wu, Haixu Wu, Zihan Qiu, Jianmin Wang, and Mingsheng Long. Supported policy optimization for offline reinforcement learning. Advances in Neural Information Processing Systems , 35:31278-31291, 2022.

Yifan Wu, George Tucker, and Ofir Nachum. The laplacian in rl: Learning representations with efficient approximations. arXiv preprint arXiv:1810.04586 , 2018.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua Susskind, Jian Zhang, Ruslan Salakhutdinov, and Hanlin Goh. Uncertainty weighted actor-critic for offline reinforcement learning. arXiv preprint arXiv:2105.08140 , 2021.

Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In Proceedings of the 39th International Conference on Machine Learning , pages 24631-24645, 17-23 Jul 2022.

Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodriguez. Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline RL. In Proceedings of the 40th International Conference on Machine Learning , volume 202, pages 38989-39007, 23-29 Jul 2023.

Rui Yang, Chenjia Bai, Xiaoteng Ma, Zhaoran Wang, Chongjie Zhang, and Lei Han. Rorl: Robust offline reinforcement learning via conservative smoothing. Advances in Neural Information Processing Systems , 35:23851-23866, 2022.

Rui Yang, Lin Yong, Xiaoteng Ma, Hao Hu, Chongjie Zhang, and Tong Zhang. What is essential for unseen goal generalization of offline goal-conditioned rl? In International Conference on Machine Learning , pages 39543-39571. PMLR, 2023.

Shentao Yang, Zhendong Wang, Huangjie Zheng, Yihao Feng, and Mingyuan Zhou. A behavior regularized implicit policy for offline reinforcement learning. arXiv preprint arXiv:2202.09673 , 2022.

Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don't change the algorithm, change the data: Exploratory data for offline reinforcement learning. arXiv preprint arXiv:2201.13425 , 2022.

Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems , 34:28954-28967, 2021.

Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems , 33:14129-14142, 2020.

Andrea Zanette, Martin J Wainwright, and Emma Brunskill. Provable benefits of actor-critic methods for offline reinforcement learning. Advances in neural information processing systems , 34:13626-13640, 2021.

Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In international conference on machine learning , pages 27042-27059. PMLR, 2022.

Appendices

A Experimental Details

A.1 ExORL Domains

We consider two locomotion and two goal-directed domains from the ExORL benchmark [95] which is built atop the DeepMind Control Suite [80]. Environments are visualised here: https: //www.youtube.com/watch?v=rAai4QzcYbs . The domains are summarised in Table 3.

Walker. A two-legged robot required to perform locomotion starting from bent-kneed position. The state and action spaces are 24 and 6-dimensional respectively, consisting of joint torques, velocities and positions. ExORL provides four tasks stand, walk, run and flip. The reward function for stand motivates straightened legs and an upright torse; walk and run are supersets of stand including reward for small and large degrees of forward velocity; and flip motivates angular velocity of the torso after standing. Rewards are dense.

Quadruped. A four-legged robot required to perform locomotion inside a 3D maze. The state and action spaces are 78 and 12-dimensional respectively, consisting of joint torques, velocities and positions. ExORL provides five tasks stand, roll, roll fast, jump and escape. The reward function for stand motivates a minimum torse height and straightened legs; roll and roll fast require the robot to flip from a position on its back with varying speed; jump adds a term motivating vertical displacement to stand; and escape requires the agent to escape from a 3D maze. Rewards are dense.

Maze. A 2D maze with four rooms where the task is to move a point-mass to one of the rooms. The state and action spaces are 4 and 2-dimensional respectively; the state space consists of x, y positions and velocities of the mass, the action space is the x, y tilt angle. ExORL provides four reaching tasks top left, top right, bottom left and bottom right. The mass is always initialised in the top left and the reward is proportional to the distance from the goal, though is sparse i.e. it only registers once the agent is reasonably close to the goal.

Jaco. A 3D robotic arm tasked with reaching an object. The state and action spaces are 55 and 6-dimensional respectively and consist of joint torques, velocities and positions. ExORL provides four reaching tasks top left, top right, bottom left and bottom right. The reward is proportional to the distance from the goal object, though is sparse i.e. it only registers once the agent is reasonably close to the goal object.

A.2 ExORL Datasets

We train on 100,000 transitions uniformly sampled from three datasets on the ExORL benchmark collected by different unsupervised agents: RANDOM, DIAYN, and RND. The state coverage on Maze is depicted in Figure 8. Though harder to visualise, we found that state marginals on higherdimensional tasks (e.g. Walker) showed a similar diversity in state coverage.

RND . An agent whose exploration is directed by the predicted error in its ensemble of dynamics models. Informally, we say RND datasets exhibit high state diversity.

DIAYN . An agent that attempts to sequentially learn a set of skills. Informally, we say DIAYN datasets exhibit medium state diversity.

RANDOM . A agent that selects actions uniformly at random from the action space. Informally, we say RANDOM datasets exhibit low state diversity.

DIAYN-100k

RND-100k

A.3 D4RL Domains

We consider two MuJoCo [81] locomotion tasks from the D4RL benchmark [21], which is built atop the v2 Open AI Gym [8]. The below environment descriptions are taken from [8].

Walker2D-v2 . A two-dimensional two-legged figure that consist of seven main body parts - a single torso at the top (with the two legs splitting after the torso), two thighs in the middle below the torso, two legs in the bottom below the thighs, and two feet attached to the legs on which the entire body rests. The goal is to walk in the in the forward (right) direction by applying torques on the six hinges connecting the seven body parts.

HalfCheetah-v2. A 2-dimensional robot consisting of 9 body parts and 8 joints connecting them (including two paws). The goal is to apply a torque on the joints to make the cheetah run forward (right) as fast as possible, with a positive reward allocated based on the distance moved forward and a negative reward allocated for moving backward.

A.4 D4RL Datasets

We consider three goal-directed datasets from D4RL, each providing a different proportion of expert trajectories. The below dataset descriptions are taken from [21].

Medium. Generated by training an SAC policy, early-stopping the training, and collecting 1M samples from this partially-trained policy.

Medium-replay. Generated by recording all samples in the replay buffer observed during training until the policy reaches the 'medium' level of performance.

Medium-expert. Generated by mixing equal amounts of expert demonstrations and suboptimal data, either from a partially trained policy or by unrolling a uniform-at-random policy.

A.5 Evaluation Protocol

We evaluate the cumulative reward (hereafter called score) achieved by VC-FB, MC-FB and our baselines on each task across five seeds. We report task scores as per the best practice recommendations of [1]. Concretely, we run each algorithm for 1 million learning steps, evaluating task scores at checkpoints of 20,000 steps. At each checkpoint, we perform 10 rollouts, record the score of each, and find the interquartile mean (IQM). We average across seeds at each checkpoint to create the learning curves reported in Appendix F. From each learning curve, we extract task scores from the learning step for which the all-task IQM is maximised across seeds. Results are reported with 95% confidence intervals obtained via stratified bootstrapping [17]. Aggregation across tasks, domains and datasets is always performed by evaluating the IQM. Full implementation details are provided in Appendix B.1.

A.6 Computational Resources

We train our models on NVIDIA A100 GPUs. Training a single-task offline RL method to solve one task on one GPU takes approximately 4 hours. FB and SF solve one domain (for all tasks) on one GPU in approximately 4 hours. Conservative FB variants solve one domain (for all tasks) on one GPU in approximately 12 hours. As a result, our core experiments on the 100k datasets used approximately 110 GPU days of compute.

B Implementation Details

Here we detail implementations for all methods discussed in this paper. The code required to reproduce our experiments is available via https://github.com/enjeeneer/zero-shot-rl .

B.1 Forward-Backward Representations

B.1.1 Architecture

The forward-backward architecture described below follows the implementation by [83] exactly, other than the batch size which we reduce from 1024 to 512. We did this to reduce the computational expense of each run without limiting performance. The hyperparameter study in Appendix J of [83] shows this choice is unlikely to affect FB performance. All other hyperparameters are reported in Table 4.

Forward Representation F ( s, a, z ) . The input to the forward representation F is always preprocessed. State-action pairs ( s, a ) and state-task pairs ( s, z ) have their own preprocessors which are feedforward MLPs that embed their inputs into a 512-dimensional space. These embeddings are concatenated and passed through a third feedforward MLP F which outputs a d -dimensional embedding vector. Note: the forward representation F is identical to ψ used by USF so their implementations are identical (see Table 4).

Backward Representation B ( s ) . The backward representation B is a feedforward MLP that takes a state as input and outputs a d -dimensional embedding vector.

Actor π ( s, z ) . Like the forward representation, the inputs to the policy network are similarly preprocessed. State-action pairs ( s, a ) and state-task pairs ( s, z ) have their own preprocessors which feedforward MLPs that embed their inputs into a 512-dimensional space. These embeddings are concatenated and passed through a third feedforward MLP which outputs a a -dimensional vector, where a is the action-space dimensionality. A Tanh activation is used on the last layer to normalise their scale. As per [23]'s recommendations, the policy is smoothed by adding Gaussian noise σ to the actions during training. Note the actors used by FB and USFs are identical (see Table 4).

Misc. Layer normalisation [3] and Tanh activations are used in the first layer of all MLPs to standardise the inputs.

B.1.2 Task Sampling Distribution Z

FB representations require a method for sampling the task vector z at each learning step. [83] employ a mix of two methods, which we replicate:

Uniform sampling of z on the hypersphere surface of radius √ d around the origin of R d ,

Biased sampling of z by passing states s ∼ D through the backward representation z = B ( s ) . This also yields vectors on the hypersphere surface due to the L 2 normalisation described above, but the distribution is non-uniform.

We sample z 50:50 from these methods at each learning step.

B.1.3 Maximum Value Approximator µ

The conservative variants of FB require access to a policy distribution µ ( a | s ) that maximises the value of the current Q iterate in expectation. Recall the standard CQL loss

where α is a scaling parameter, µ ( a | s ) the policy distribution we seek, R regularises µ and L Q represents the normal TD loss on Q . [44]'s most performant CQL variant (CQL( H )) utilises maximum entropy regularisation on µ i.e. R = H ( µ ) . They show that obtaining µ can be cast as a closed-form optimisation problem of the form:

and has optimal solution µ ∗ ( x ) = 1 Z exp( f ( x )) , where Z is a normalising factor. Plugging Equation 14 into Equation 13 we obtain:

In discrete action spaces the logsumexp can be computed exactly; in continuous action spaces [44] approximate it via importance sampling using actions sampled uniformly at random, actions from the current policy conditioned on s t ∼ D , and from the current policy conditioned on s t +1 ∼ D 2 :

with N a hyperparameter defining the number of actions to sample across the action-space. We can substitute F ( s, a, z ) ⊤ z for Q ( s, a ) in the final expression of Equation 17 to obtain the equivalent for VC-FB:

In Appendix F, Figure 16 we show how the performance of VC-FB varies with the number of action samples. In general, performance improves with the number of action samples, but we limit N = 3 to limit computational burden. The formulation for MC-FB is identical other than each value F ( s, a, z ) T z being replaced with measures F ( s, a, z ) T B ( s + ) .

B.1.4 Dynamically Tuning α

A critical hyperparameter is α which weights the conservative penalty with respect to other losses during each update. We initially trialled constant values of α , but found performance to be fragile to this selection, and lacking robustness across environments. Instead, we follow [44] once again, and instantiate their algorithm for dynamically tuning α , which they call Lagrangian dual gradient-descent on α . We introduce a conservative budget parameterised by τ , and set α with respect to this budget:

Intuitively, this implies that if the scale of overestimation ≤ τ then α is set close to 0, and the conservative penalty does not affect the updates. If the scale of overestimation ≥ τ then α is set proportionally to this gap, and thus the conservative penalty is proportional to the degree of overestimation above τ . As above, for the MC-FB variant values F ( s, a, z ) ⊤ z are replaced with measures F ( s, a, z ) ⊤ B ( s + ) .

B.1.5 Algorithm

We summarise the end-to-end implementation of VC-FB as pseudo-code in Algorithm 1. MC-FB representations are trained identically other than at line 10 where the conservative penalty is computed for M instead of Q , and in line 12 where M s are lower bounded via Equation 12.

B.1.6 Directed Value-Conservative Forward Backward Representations

VC-FB applies conservative updates using task vectors z sampled from Z (which in practice is a uniform distribution over the √ d -hypersphere). This will include many vectors corresponding to tasks that are never evaluated in practice in downstream applications. Intuitively, it may seem reasonable to direct conservative updates to focus on tasks that are likely to be encountered downstream. One simple way of doing this would be consider the set of all goal-reaching tasks for goal states in the training distribution, which corresponds to sampling z = B ( s g ) for some s g ∼ D . This leads to the following conservative loss function:

We call models learnt via this loss directed -VC-FB ( D VC-FB). While we were initially open to the possibility that D VC-FB updates would be better targeted than those of VC-FB, and would lead to improved downstream task performance, this turns out not to be the case in our experimental settings as discussed in Section 5. We report scores obtained by the D VC-FB method across all 100k datasets, domains and tasks in Appendix C.

B.2 Universal Successor Features

We directly reimplement USFs, with basic features φ ( s ) provided by Laplacian eigenfunctions [87], from [83].

B.2.1 Architecture

USF ψ ( s, a, z ) . The input to the USF ψ is always preprocessed. State-action pairs ( s, a ) and state-task pairs ( s, z ) have their own preprocessors which are feedforward MLPs that embed their inputs into a 512-dimensional space. These embeddings are concatenated and passed through a third

feedforward MLP ψ which outputs a d -dimensional embedding vector. Note this is identical to the implementation of F as described in Appendix B.1. All other hyperparameters are reported in Table 4.

Feature Embedding φ ( s ) . The feature map φ ( s ) is a feedforward MLP that takes a state as input and outputs a d -dimensional embedding vector. The loss function for learning the feature embedding is provided in Appendix B.2.2.

Actor π ( s, z ) . Like the forward representation, the inputs to the policy network are similarly preprocessed. State-action pairs ( s, a ) and state-task pairs ( s, z ) have their own preprocessors which are feedforward MLPs that embed their inputs into a 512-dimensional space. These embeddings are concatenated and passed through a third feedforward MLP which outputs a a -dimensional vector, where a is the action-space dimensionality. A Tanh activation is used on the last layer to normalise their scale. As per [23]'s recommendations, the policy is smoothed by adding Gaussian noise σ to the actions during training. Note this is identical to the implementation of π ( s, z ) as described in Appendix B.1.

Misc. Layer normalisation [3] and Tanh activations are used in the first layer of all MLPs to standardise the inputs. z sampling distribution Z is identical to FB's (Appendix B.1.2).

B.2.2 Laplacian Eigenfunctions Loss

Laplacian eigenfunction features φ ( s ) are learned as per [87]. They consider the symmetrized MDP graph Laplacian created by some policy π , defined as L = Id -1 2 ( P π diag ρ -1 + diag ρ -1( P π ) T ) . They learn the eigenfunctions of L with the following:

which comes from [39].

B.3 GC-IQL

B.3.1 Architecture

We implement GC-IQL following [60]'s codebase. GC-IQL inherits all functionality from a base soft actor-critic agent [28], but adds a soft conservative penalty to the goal-conditioned critic's V ( s, g ) updates. We refer the reader to paper that introduces GC-IQL [60] for details on the loss function used to train V ( s, g ) . Hyperparameters are reported in Table 5.

Critic(s). GC-IQL trains double goal-conditioned value functions V ( s, g ) . The critics are feedforward MLPs that take a state-goal pair ( s, g ) as input and output a value ∈ R 1 . During training the goals are sampled from the prior G described in Section B.3.2.

Actor. The actor is a standard feedforward MLP taking the state s as input and outputting an 2 a -dimensional vector, where a is the action-space dimensionality. The actor predicts the mean and standard deviation of a Gaussian distribution for each action dimension; during training a value is sampled at random, during evaluation the mean is used.

B.3.2 Goal Sampling Distribution G

Following [60], goals are sampled from either random states, future states, or the current state with probabilities 0 . 3 , 0 . 5 and 0 . 2 respectively. A geometric distribution Geom (1 -γ ) is used for the future state distribution, and the uniform distribution over the offline dataset is used for sampling random states.

B.4 CQL

B.4.1 Architecture

We adopt the same implementation and hyperparameters as is used on the ExORL benchmark. CQL inherits all functionality from a base soft actor-critic agent [28], but adds a conservative penalty to the critic updates (Equation 10). Hyperparameters are reported in Table 5.

Critic(s). CQL employs double Q networks, where the target network is updated with Polyak averaging via a momentum coefficient. The critics are feedforward MLPs that take a state-action pair ( s, a ) as input and output a value ∈ R 1 .

Actor. The actor is a standard feedforward MLP taking the state s as input and outputting an 2 a -dimensional vector, where a is the action-space dimensionality. The actor predicts the mean and standard deviation of a Gaussian distribution for each action dimension; during training a value is sampled at random, during evaluation the mean is used.

B.5 TD3

B.5.1 Architecture

We adopt the same implementation and hyperparameters as is used on the ExORL benchmark. Hyperparameters are reported in Table 5.

Critic(s). TD3 employs double Q networks, where the target network is updated with Polyak averaging via a momentum coefficient. The critics are feedforward MLPs that take a state-action pair ( s, a ) as input and output a value ∈ R 1 .

Actor. The actor is a standard feedforward MLP taking the state s as input and outputting an a -dimensional vector, where a is the action-space dimensionality. The policy is smoothed by adding Gaussian noise σ to the actions during training.

Misc. As is usual with TD3, layer normalisation [3] is applied to the inputs of all networks.

B.6 Code References

This work was enabled by: NumPy [30], PyTorch [61], Pandas [56] and Matplotlib [31].

C Extended Results

In this section we report a full breakdown of our experimental results on the ExORL benchmark by dataset, domain and task. Table 6 reports results for methods trained on the 100k sub-sampled datasets, and Table 7 reports results for methods trained on the full datasets.

D Value Conservative Universal Successor Features

In this section, we develop value conservative regularisation for use by Universal Successor Features (USF) [5, 7], the primary alternative to FB for zero-shot RL.

Recall from Section 2 that successor features require a state-feature mapping φ : S → R d which is usually obtained by some representation learning method [5]. Universal successor features are the expected discounted sum of these features, starting in state s 0 , taking action a 0 and following the task-dependent policy π z thereafter

USFs satisfy a Bellman equation [7] and so can be trained using TD-learning on the Bellman residuals:

where ¯ ψ is a lagging target network updated via Polyak averaging, and Z is identical to that used for FB training (Appendix B.1.2). As with FB representations, the policy maximises the Q function defined by ψ :

and for continuous state and action spaces is trained in an actor critic formulation. Like FB, USF training requires next action samples a t +1 ∼ π z ( s t +1 ) for the TD targets. We therefore expect SFs to suffer the same failure mode discussed in Section 3 (OOD state-action value overestimation) and to benefit from the same remedial measures (value conservatism). Training value-conservative successor features (VC-SF) amounts to substituting the USF Q function definition and loss for FB's in Equation 11:

Both the maximum value approximator µ ( a | s ) (Equation 17, Section B.1.3) and α -tuning (Equation 18, Section B.1.4) can be extracted identically to the FB case with any occurrence of F ( s, a, z ) ⊤ z substituted with ψ ( s, a, z ) ⊤ z . As USFs do not predict successor measures we cannot formulate measure-conservative USFs.

E Negative Results

In this section we provide detail on experiments we attempted, but which did not provide results significant enough to be included in the main body.

E.1 Downstream Finetuning

If we relax the zero-shot requirement, could pre-trained conservative FB representations be finetuned on new tasks or domains? Base CQL models have been finetuned effectively on unseen tasks using both online and offline data [41], and we had hoped to replicate similar results with VC-FB and MC-FB. We ran offline and online finetuning experiments and provide details on their setups and results below. All experiments were conducted on the Walker domain.

Offline finetuning. We considered a setting where models are trained on a low quality dataset initially, before a high quality dataset becomes available downstream. We used models trained on the RANDOM-100k dataset and finetuned them on both the full RND and RND-100k datasets, with models trained from scratch used as our baseline. Finetuning involved the usual training protocol as described in Algorithm 1, but we limited the number of learning steps to 250k.

We found that though performance improved during finetuning, it improved no quicker than the models trained from scratch. This held for both the full RND and RND-100k datasets. We conclude that the parameter initialisation delivered after training on a low quality dataset does not obviously expedite learning when high quality data becomes available.

Online finetuning. We considered the online finetuning setup where a trained representation is deployed in the target environment, required to complete a specified task, and allowed to collect a replay buffer of reward-labelled online experience. We followed a standard online RL protocol where a batch of transitions was sampled from the online replay buffer after each environment step for use in updating the model's parameters. We experimented with fixing z to the target task during in the actor updates (Line 16, Algorithm 1), but found it caused a quick, irrecoverable collapse in actor performance. This suggested uniform samples from Z provide a form of regularisation. We granted the agents 500k steps of interaction for online finetuning.

We found that performance never improved beyond the pre-trained (init) performance during finetuning. We speculated that this was similar to the well-documented failure mode of online finetuning of CQL [59], namely taking sub-optimal actions in the real env, observing unexpectedly high reward, and updating their policy toward these sub-optimal actions. But we note that FB representations do not update w.r.t observed rewards, and so conclude this cannot be the failure mode. Instead it seems likely that FB algorithms cannot

FB

MC-FB

VC-FB (scratch)

VC-FB

■FB (scratch)

MC-FB (scratch)

350

300

250

(WoI)

200

Return

150

100

50

0

0

50

100

150

200

250

Learning steps (k)

use the narrow, unexploratory experience obtained from attempting to perform a specific task to improve model performance.

We believe resolving issues associated with finetuning conservative FB algorithms once the zero-shot requirement is relaxed is an important future direction and hope that details of our negative attempts to this end help facilitate future research.

FB

MC-FB

VC-FB (init)

VC-FB

FB (init)

MC-FB (init)

600

500

(IQM)

400

Return (

300

200

100

0

50

100

150

200

250

Learning steps (k)

F Learning Curves & Hyperparameter Sensitivity

ExORL Learning Curves. We report the learning curves for all zero-shot RL methods and CQL in Figures 11, 12, and 13. For all domains except Jaco, the y-axis limit is fixed at 1000 as that is the maximum score achievable in the DeepMind Control Suite. For Jaco-related figures, the y-axis limits is fixed at 100 as no method achieves a score higher than this.

Hyperparameter Sensitivity. We report the sensitivity of VC-FB and MC-FB to the choice of two new hyperparameters: conservative budget τ and action samples per policy N on the ExORL benchmark. Figure 14 plots the sensitivity of VC-FB to the choice of τ on Walker and Maze domains across RND and RANDOM datasets. Figure 15 plots the sensitivity of MC-FB to the choice of τ on Walker and Maze domains across RND and RANDOM datasets. Figure 16 plots the sensitivity of MC-FB to the choice of N on Walker and Maze domains across RND and RANDOM datasets.

We further explore the sensitivity of VC-FB performance on Walker2D from the D4RL benchmark w.r.t. the choice of conservative budget τ . Figure 17 plots this relationship when trained on the 'medium-expert" dataset from D4RL.

walker-all-tasks-rnd

walker-flip-rnd

walker-run-rnd

walker-stand-rnd

1000

1000

1000

1000

(WOI)

800

800

800

800

600

600

600

600

Return

400

400

400

400

200

200

200

200

W

0

0

0

0

walker-walk-rnd

walker-all-tasks-random

walker-flip-random

walker-run-random

1000

1000

1000

1000

(WOI)

800

800

800

800

600

600

600

600

Return(

400

400

400

400

200

200

200

200

0

0

0

0

walker-stand-random

walker-walk-random

walker-all-tasks-diayn

walker-flip-diayn

1000

1000

1000

1000

(WoI)

800

800

800

800

600

600

600

600

Return

400

400

400

400

200

200

200

ARA

200

MARMMT

0

0

0

0

walker-run-diayn

walker-stand-diayn

walker-walk-diayn

maze-all-tasks-rnd

1000

1000

1000

1000

(IQM)

800

800

800

800

600

600

600

600

Return

400

400

400

400

200

200

200

200

MV

0

0

0

0

maze-btm-left-rnd

maze-btm-right-rnd

maze-top-left-rnd

maze-top-right-rnd

1000

1000

1000

1000

(Wol)

800

800

800

800

600

600

600

600

Return

400

400

400

400

200

200

200

200

0

0

0

0

maze-all-tasks-random

maze-btm-left-random

maze-btm-right-random

maze-top-left-random

1000

1000

1000

1000

(WOl)

800

800

800

一

800

600

600

600

600

Return(

400

400

400

400

200

200

200

200

0

0

0

0

0.0 0.20.4 0.6 0.81.0

0.00.20.4 0.60.81.0

0.00.20.40.6 0.81.0

0.00.20.40.60.81.0

Learning steps (×106)

Learning steps (×106)

Learning steps (×106)

Learning steps (×106)

SF-LAP

GC-IQL

FB

VC-FB

MC-FB

■CQL

maze-top-right-random

maze-all-tasks-diayn

maze-btm-left-diayn

maze-btm-right-diayn

1000

1000

1000

1000

(WOI)

800

800

800

800

600

600

600

600

Return

400

400

400

400

200

200

200

200

0

0

0

0

maze-top-left-diayn

maze-top-right-diayn

jaco-all-tasks-rnd

jaco-btm-left-rnd

1000

1000

100

100

Return (IQM)

800

800

80

80

600

600

60

60

400

400

40

40

200

200

20

20

0

0

0

0

jaco-btm-right-rnd

jaco-top-left-rnd

jaco-top-right-rnd

jaco-all-tasks-random

100

100

100

100

Return (IQM)

80

80

80

80

60

60

60

60

40

40

40

40

20

20

20

20

0

0

0

0

jaco-btm-left-random

jaco-btm-right-random

jaco-top-left-random

jaco-top-right-random

100

100

100

100

Return (IQM)

80

80

80

80

60

09

60

60

40

40

40

40

20

20

一

20

20

0

0

0

0

jaco-all-tasks-diayn

jaco-btm-left-diayn

jaco-btm-right-diayn

jaco-top-left-diayn

100

100

100

100

(Wol)

80

80

80

---------

80

60

60

60

= =

60

Return(

40

40

40

40

20

20

20

20

0

0

0

0

jaco-top-right-diayn

quadruped-all-tasks-rnd

quadruped-escape-rnd

puu-dwn!-padnupenb

100

500

500

500

Return (IQM)

80

400

400

400

60

300

300

300

40

200

200

200

20

100·

100

100

0

0

0

0

0.0 0.20.4 0.6 0.81.0

0.0 0.20.4 0.6 0.81.0

0.0 0.20.40.6 0.81.0

0.0 0.20.4 0.6 0.81.0

Learning steps (×106)

Learning steps (×106)

Learning steps (×106)

Learning steps (×106)

SF-LAP

GC-IQL

FB

VC-FB

MC-FB

■CQL

quadruped-roll-rnd

quadruped-roll-fast-rnd

quadruped-stand-rnd quadruped-all-tasks-random

500

500

500

500

(IQM)

400

400

400

400

300

300

300

300

Return（

200

200

200

200

100

100

100

100

WA

0

0

0

0

quadruped-escape-randomquadruped-jump-random quadruped-roll-random quadruped-roll-fast-random

500

500

500

500

(IQM)

400

400

400

400

300

300

300

300

Return

200

200

200

200

100

100

100

100

0

0

0

0

uei-duno-padnpenbueip-adeosa-padnpenb ureip-sysea-lle-padnupenbwopues-puess-padnupenb

500

1000

1000

1000

(IQM)

400

800

800

800

300

600

600

600

Return（

200

400

400

400

100

200

200

200

0

0

0

0

0.00.20.40.60.81.0

quadruped-roll-diayn

quadruped-roll-fast-diaynquadruped-stand-diayn

Learning steps (×106)

1000

1000

1000

(WoI)

800

800

800

600

600

600

Return

400

400

400

200

200

200

0

0

0

0.00.20.40.6 0.81.0

0.00.20.4 0.60.81.0

0.0 0.20.40.6 0.81.0

Learning steps (×106)

Learning steps (×106)

Learning steps (×106)

SF-LAP

GC-IQL

FB

VC-FB

MC-FB

CQL

walker-rnd

maze-rnd

600

600

(WOl)

400

400

Return

200

200

0

0

10

25

50

75

100

10

25

50

75

100

walker-random

maze-random

200

100

(WOI)

150

75

Return

100

50

50

25

0

1025

50

75

100

10

25

50

75

100

Conservative budgetT

Conservative budgetT

walker-rnd

maze-rnd

600

600

(Wol)

400

400

Return(

200

200

0

0

25

50

75

100

25

50

75

100

walker-random

maze-random

200

100

(Wo)

150

75

Return

100

50

50

25

-0

一

0

25

50

75

100

25

50

75

100

Conservative budget T

Conservative budget T

walker-rnd

maze-rnd

400

(Wol)

400

Return

200

200

0

0

1

3

4

1

3

4

walker-random

maze-random

200

100

(IQM)

150

75

Return

100

50

50

25

0

0

1

3

4

1

3

4

Action samples per policy N

Action samples per policy N

100

80

60

40

20

0

25

35404550

75

Conservative budget T

G Code Snippets

G.1 Update Step

G.2 Value-Conservative Penalty

G.3 Measure-Conservative Penalty

G.4 α Tuning

H NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: This paper does not include theoretical proofs.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Code and hyperparameters are provided; datasets and environments are open-source. Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

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: See Appendix A and B, and https://github.com/enjeeneer/zero-shot-rl .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/public/ guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

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: See Appendix A and B.

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.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

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.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: See Appendix A.3.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification:

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper discusses methods for improving the performance of zero-shot RL methods when trained on low quality offline datasets. We do not expect this work to directly impact society positively or negatively.

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

11. 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: This paper poses no such risks.

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:

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.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: https://github.com/enjeeneer/zero-shot-rl .

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification:

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe work investigates methods for zero-shot reinforcement learning (RL) that can be trained on small, homogeneous datasets. This research is driven by the need to make zero-shot RL practical when large, heterogeneous datasets are unavailable. The authors identify the limitations of existing methods that overestimate out-of-distribution (OOD) state-action values when trained on low-quality datasets. They propose incorporating conservatism into zero-shot RL algorithms to mitigate these issues. Their experimental results demonstrate that conservative zero-shot RL methods outperform their non-conservative counterparts on low-quality datasets while maintaining competitive performance on high-quality datasets.\n\nSTRENGTHS:\n* Addressing Significant Issue: The paper attempts to address a significant gap in zero-shot RL by focusing on the challenges of using small, homogeneous datasets, which are more common in real-world applications.\n\n* Improved Performance: The proposed conservative zero-shot RL methods consistently outperform non-conservative counterparts on \nlow-quality datasets and do not degrade performance on high-quality datasets.\n\nWEAKNESSES:\n1. **The problem definition**: I do not think the settings of this work is so-called Zero-shot RL since the work introduces another dataset. And the code link is invalid.\n2. **Complexity of Implementation**: The introduction of conservatism increases the complexity of the algorithms, which may pose implementation challenges for practitioners.\n3. **Limited Real-World Validation**: While the experiments are thorough, there is limited validation of the proposed methods in real-world settings, which could affect their applicability.\n4. **Comparative Analysis**: The paper could benefit from a more in-depth comparative analysis with other state-of-the-art methods to highlight specific strengths and weaknesses.\n\nQUESTIONS:\nSee in weakness.\n\nLIMITATIONS:\nSee in weakness.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work addresses zero-shot reinforcement learning (RL), focusing on training agents to perform tasks without explicit rewards during pre-training. The authors investigate the performance degradation that occurs with small, homogeneous datasets and propose conservative zero-shot RL algorithms inspired by single-task offline RL approaches. Experimental results demonstrate that the proposed methods improve performance on sub-optimal datasets without compromising effectiveness on large, diverse datasets.\n\nSTRENGTHS:\n* Zero-shot pretraining for RL is a crucial problem that can enhance generalization in downstream tasks.\n\n* The proposed method is well-motivated, and the writing is clear.\n\n* The experiments are thorough and effectively validate the proposed method.\n\nWEAKNESSES:\n* One weakness of the paper is the lack of a clear definition for \"low-quality data.\" As shown in Figure 8, it might be more accurately described as \"coverage\" rather than quality. Providing a more explicit definition is recommended.\n\n* Additionally, the comparison of baselines is incomplete, as it lacks an offline goal-conditioned baseline, which is highly relevant in the considered tasks. Including a representative offline GCRL baseline, such as GC-IQL, would strengthen the comparison.\n\n* Some related works are missing, particularly in the offline GCRL domain. References [1][2][3][4] are pertinent as they also focus on learning from offline datasets without rewards and deploying for multiple goals. The didactic example in this paper is very similar to Figure 1 in a prior paper [3], which also convey similar information during zero-shot deployment for multiple goals. A discussion with the prior results would be beneficial. Furthermore, a very relevant offline zero-shot RL method [5] is neither cited nor discussed. \n\n\n*References*\n\n[1] Yang R, Lu Y, Li W, et al. Rethinking goal-conditioned supervised learning and its connection to offline rl. ICLR, 2022.\n\n[2] Park S, Ghosh D, Eysenbach B, et al. Hiql: Offline goal-conditioned rl with latent states as actions. Advances in Neural Information Processing Systems, 2024.\n\n[3] Yang R, Yong L, Ma X, et al. What is essential for unseen goal generalization of offline goal-conditioned rl?[C]//International Conference on Machine Learning. PMLR, 2023.\n\n[4] Eysenbach B, Zhang T, Levine S, et al. Contrastive learning as goal-conditioned reinforcement learning[J]. Advances in Neural Information Processing Systems, 2022.\n\n[5] Frans K, Park S, Abbeel P, et al. Unsupervised Zero-Shot Reinforcement Learning via Functional Reward Encodings[J]. arXiv preprint arXiv:2402.17135, 2024.\n\nQUESTIONS:\n* Could the authors provide a more explicit definition of \"low-quality data\"?\n\n* It is recommended to compare the proposed method with a representative offline GCRL baseline.\n\n* Including related works discussed above would make this study more comprehensive.\n\n* Can this method also be applied to approaches based on successor features?\n\nLIMITATIONS:\nOne limitation is that VC-FB and MC-FB do not outperform CQL on the D4RL benchmark; addressing this issue is left for future work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper identifies an overestimation problem of zero-shot RL algorithms, particularly in low data or low data quality settings. To address this, they propose to use a value conservative method to mitigate the overestimation (from CQL). They showcase that this effectively allows for reducing the overestimation, then empirically show performance improvements on low or bad quality data regimes, without impacting performance on usual regimes.

STRENGTHS:
I am not an expert of zero short RL, but the paper seems sound and is mostly clearly written. The problem seems relevant as zero short RL from low quality data is relevant for robotic applications, as underlined by the authors. I like the first figure, that provides intuition, similar to the ones in DDQN's paper. The method is clearly motivated and explained, the experimental evaluation is clear, the figure and caption allow assessing the defined precise scientific questions. I do not see much flaw in this paper presentation and method, but my qualification is not too high.

WEAKNESSES:
**Provide a bit more background and intuition about zero-shot RL**. This is not really mandatory, but it would help a broader audience.

**Introduce the RANDOM and RND datasets**. Maybe an introduction of these datasets can help better grasp Figure 2 and 3, and would help the previous point.

QUESTIONS:
* What is the intuition behind the overestimation reduction of CQL ?

LIMITATIONS:
There is a dedicated limitation section that correctly addresses the limitations of the method.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

comments: SUMMARY:
This work proposes modifications to Forward-Backward representations method. The modifications are aimed to improve the method's robustness to dataset quality. The vanilla FB suffers when the dataset does not cover the state space well, and overestimates the quality of actions. The authors show this with a simple example on point-mass. Two modifications are proposed, VC-FB and MC-FB, which take inspiration from CQL and make sure the FB representations are not too optimistic on the transitions not seen in the dataset. 
VC-FB and MC-FB are then thoroughly tested on ExORL and D4RL. The authors show that VC-FB offers the best improvement on ExORL, although none of the proposed modifications beat CQL on D4RL benchmark. Nevertheless, in all settings the modifications greatly improve over vanilla FB.

STRENGTHS:
- The paper is clearly written.
- The proposed modifications have clear mathematical motivation, similar to CQL
- Both the modifications improve over the vanilla FB representation, and even outperform CQL on ExORL
- The experiments are thorough and clear

WEAKNESSES:
- The novelty is limited: the authors take CQL's method for correcting overly optimal predictions and apply them to FB.

QUESTIONS:
-  How do the proposed modifications affect the performance of the model when the dataset is of good quality? Does it hurt the performance?

LIMITATIONS:
The authors have adequately described limitations and potential negative societal impact of their work.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12