Provable Zero-Shot Generalization in Offline Reinforcement Learning
Authors: Zhiyong Wang
Chen Yang
John C.S. Lui
Dongruo Zhou
Venue: ICML 2025 poster (2025)
Keywords: offline reinforcement learning
generalization
Primary area: theory->reinforcement_learning_and_planning

Abstract: In this work, we study offline reinforcement learning (RL) with zero-shot generalization property (ZSG), where the agent has access to an offline dataset including experiences from different environments, and the goal of the agent is to train a policy over the training environments which performs well on test environments without further interaction. Existing work showed that classical offline RL fails to generalize to new, unseen environments. We propose pessimistic empirical risk minimization (PERM) and pessimistic proximal policy optimization (PPPO), which leverage pessimistic policy evaluation to guide policy learning and enhance generalization. We show that both PERM and PPPO are capable of finding a near-optimal policy with ZSG. Our result serves as a first step in understanding the foundation of the generalization phenomenon in offline reinforcement learning.

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Zhiyong Wang'
Chen Yang²
John C.S. Lui' Dongruo Zhou?
Abstract
been studied in online RL Rajeswaran al. 2017 Machado et al. 2018 Justesen et al., 2018 Packer et al. 2019 Zhang
In this work, we study offline reinforcement learn ing (RL) with zero-shot generalization property (ZSG), where the agent has access to an offline dataset including experiences from different envi- ronments, and the goal of the agent is to train policy over the training environments which per- forms well on test environments without further interaction. Existing work showed that classical offline RI. fails to generalize to new, unseen envi- ronments. We propose pessimistic empirical risk minimization (PERM) and pessimistic proximal policy optimization (PPPO), which leverage pes- simistic policy evaluation to guide policy learning and enhance generalization. We show that both PERM and PPPO are capable of finding a near- optimal policy with ZSG. Our result serves as a first step in understanding the foundation of the generalization phenomenon in offlin reinforce- ment learning.
et al. 2018 8a b), as the agent receives no training data from the environments it is tested on.
A number of recent empirical studies (Medraata et al. 2023 Yang et al. 2023 Mazoure et al., 2022) have recognized this challenge and introduced various offline RL methodolo- gies that are capable of ZSG. Notwithstanding the lack of theoretical backing, these methods are somewhat restrictive; for instance, some are only effective for environments that vary solely in observations(Mazouuu et al., 2022). while others are confined to the realm of imitation learning Yang et al. 2023), thus limiting their applicability to compre- hensive framework of offline RL with ZSG capabilities. Concurrently, theoretical advancements (Bose al. 2024] Ishfaq et al. 2024) in this domain have explored multi-task offiin RL by focusing representation learning. These approaches endeavor to derive a low-rank representation of states and actions, which inherently requires additional interactions with the downstream tasks to effectively formu- late policies based on these representations. Therefore, we raise natural question:
1. Introduction
Can we design provable offline RL with zero-shot generalization ability?
Offline reinforcement learning (RL) has become increas- ingly significant in modern RL because it eliminates the need for direct interaction between the agent and the en- vironment; instead, it relies solely on learning from an of- fline training dataset. However, in practical applications, the offline training dataset often originates from different environment than the one of interest. This discrepancy ne- cessitates evaluating RL agents in a generalization setting, where the training involves finite number of environments drawn from a specific distribution, and the testing is con- ducted on a distinct set of environments from the same or different distribution. This scenario is commonly referred to as the zero-shot generalization (ZSG) challenge which has
We propose novel offline RL frameworks that achieve ZSG to address this question affirmatively. Our contributions are listed as follows.
We first analyze when existing offline RL approaches fail to generalize without further algorithm modifications. Specifically, we prove that if the offline dataset does not contain context information, then it is impossible for vanilla RL that equips a Markovian policy to achieve ZSG property. We show that the offline dataset from contextual Markov Decision Process (MDP) is not dis- tinguishable from a vanilla MDP which is the average of contextual Markov Decision Process over all contexts. Such an analysis verifies the necessity of new RL methods with ZSG property.
'Department of Computer Science and Engineering, The Chinese University of Hong Kong, Hong Kong SAR, China Departmentt of Computer Science, Indiana University Bloom- ington, Bloomington, USA. Correspondence to: Dongruo Zhou <dz1 @iiu.duu..
We propose two meta-algorithms called pessimistic empir- ical risk minimization (PERM) and pessimistic proximal policy optimization (PPPO) that enable ZSG for offline
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
RL (Jin al. D. In detail, both of our algorithms take pessimistic policy evaluation (PPE) oracle its component and output policies based on offline datasets from multiple environments. Our result shows that the sub-optimalities of the output policies are bounded by both the supervised learning error, which is controlled by the number of different environments, and the reinforce- ment learning error, which is controlled by the coverage of the offline dataset to the optimal policy. Please refer to Table for a summary of our results. To the best of our knowledge, our proposed algorithms are the first offline RL methods that provably enjoy the ZSG property.
offiine RL methodologies. This includes model-based ap-
proaches [Rashidinejad et al. 2021 Uehara and Sun 2021
Jin et al.
2021 Yu et al., 2020; Xie et al. 202 lb Uchara
ct al.
2021
Yin et al. 2022). model-free techniques Kumar
et al.
2020
Wu et al. 2021: Bai et al. 2022 Ghasemipour
ct al. 2022
Yan et al., 2023), and policy- strategies
Rezaeifar et al. 2022, Xie et al.
2021a Zanette et al.
2021 Nguyen-Tang and Arora 2024). Yarats et al., [2022]
has observed that with sufficient offline data diversity and coverage, the need for pessimism to mitigate extrapolation errors and distribution shift might be reduced. To the best of our knowledge, we are the first to theoretically study the generalization ability of offline RL in the contextual MDP setting.
Notation We use lower case letters to denote scalars, and use lower and upper case bold face letters to denote vectors and matrices respectively. We denote by [ni] the set (1,-... n). For a vector x E I2d and positive semi-definite matrix E E Rdxd, we denote by ||||| the vector's Euclidean norm and define ||x| } and {b.,} with n = 1,2,.... we write u,, = Obr)) if there exists an absolute constant C' > 0 such that On S Ob, holds for all n 2 and write a = (2(b,)) if there exists an absolute constant C' > 0 such that a, 2 Cbn holds for all n 2 1. We use Ō() to further hide the polylogarithmic factors. We use (24)²=1 to denote sequence (I ..., In), and we use {ax 1 to denote the set (21,...,, }. We use (L((llq) to denote the KL. distance between distributions p and g. defined as f p log(p/)). 1). We use B[x)) V(z:] to denote expectation and variance of a random variable x.
Generalization in online RL There are extensive empirical studies on training online RL agents that can generalize
to new transition and reward functions (Rajeswaran et al.
2017
Machado et al.
2018
Justesen et al.
2018
Packer
Vx Ex. For two positive sequences
et al.
2019:
Zhang et al.
2018a
b Nichol et al.
2018
Cobbe
ct al. 2018
Küttler et al.
2020
Bengio et al.
2020
Bertran
et al. 2020
Ghosh et al.
2021 Kirk et al.
2023 Juliani
et al.
201 9
Ajay et al.
2021
Jan et al.
2021 Frans
and Isola
2022
Albrecht et al.) 2022
Ehrenberg al.
2022 Song et al.
2020 Lyle ct al. 2022 Ye et al. 2020
Lee et al. [0200] Jiang et al.) They use techniques including
implicit regularization (Song et al. (2020). data augmenta- tion (Ye et al. 2020 Lee et al. (020), uncertainty-drivve exploration (Jiang et al.). successor feature (Touati et al., 2023), etc. These works focus mostly on the online RL set- ting and do not provide theoretical guarantees, thus differing lot from ours. Moreover, (Touati et al.. 2023) has studied zero-shot generalization in offline RL, but to unseen reward functions rather than unseen environments. Addtional re- lated works that have studied zero-shot RI. include (Park
The remaining parts are organized as follows. In Section 2 we discuss related works. In Section we introduce the setting of our work. In Section 4 we analyze when existing offline RL approaches (Jin et al. 202T) fail to generalize without further algorithm modifications. In Section 5 we introduce our proposed meta-algorithms and provide their theoretical guarantees. In Section 6 we specify our meta- algorithms and analysis to a more concrete lincar MDP setting. Finally, in Section 7 we conclude our work and propose some future directions.
et al. 2024 Jeen et al. 2023).
There are also some recent works aimed at understanding on- line RL generalization from theoretical perspective. |aang] et al. (2019) examined specific class of reparameteriz- able RL problems and derived generalization bounds using Rademacher complexity and the PAC-Bayes bound. Ma-
lik et al.] (2021) established lower bounds and introduced efficient algorithms that ensure a near-optimal policy for deterministic MDPs. A recent work (Ye et al. (023) studied how much pre-training can improve online RL test perfor- mance under different generalization settings. To the best of our knowledge, no previous work exists on theoretical understanding of the zero-shot generalization of offline RL.
2. Related works
Offline RL Offline reinforcement learning (RL) EErstt et al.
2005
Riedmiller 2005 Lange et al., 2012 Levine et al.
2020
Wang et al. 2024) addresses the challenge of learning
a policy from pre-collected dataset without direct online interactions with the environment. A central issue in offline RL is the inadequate dataset coverage, stemming from lack of exploration (Levine et al., 2020 Liu et al. 2020). A common strategy to address this issue is the application of the pessimism principle, which penalizes the estimated value of under-covered state-action pairs. Numerous studies have integrated pessimism into various single-environment
Our paper is also related to recent works studying multi- task learning in reinforcement learning (RL) (Brunskill and
Li 2013 Tirinzoni et al., 2020 Hu et al.
2021 Zhang
and Wang 2021 Lu al.
2021 Bose et al.
2024 Ishfaq
et al. 2024 Zhang et al. 2023) which focus transferring
2

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Table 1. Summmary of our algorithms and their suhoptimality gaps, where A the action space, H is the length of episode, is the number of environments in the offline dataset. Note that in the multi-environment setting. '' is the near-optimal policy w.r.t. expectation (defined
Section 3 N is the covering number of the policy space 11 w.r.t distance d(r' , xx) = S, [m]] (-Is) - * (+1)]]]. The
uncertainty quantifier is are tailored with the oracle return in the corresponding algorithms (details are in Section Algorithm
Vlog(N)/n n Suboptimality Gap
PERM (our Algo 2
2x E [T (1,( 04)
-H
81 201
PPPO (our Algo
Vlog [A| H² f/2/nn + n = h=1
Bi,
Te, aA)
81 x1
the knowledge learned from upstream tasks to downstream ones. Additionally, these works typically assume that all tasks share similar transition dynamics or common represen- tations while we do not. Meanwhile, they typically require the agent to interact with the downstream tasks, which does not fall into the ZSG regime.
sampling of context e from the context space. Then our goal is to develop the generalizable RL with small zero-shot generalization gap (ZSG gap), defined as follows:
SubOpt(T) == Eewe [Vai (x11) - Eewe [V2 (m1)).
Remark 1 We briefty compare generalizable RL with sev- eral related settings. Robust RL (Pinto et al.) 2017) aims to find the best policy for the worst-case environment, whereas generalizable RL seeks a policy that performs well in the average-case environment. Meta-RL (Beck et al. (023) enables few-shot adaptation to new environments, either through policy updates (Finn et al. 2017) or via history- dependent policies (Duan al. 2016). In contrast, gener- alizable RL. primarily focuses the zero-shot selting. In the general POMDP framework (Cassandra et al.) 1994). agents need to maintain history policies to im- plicitly infer environment information, while generalizable RI. aims to discover a single state stpeedeeendent policy that generalizes well across all environments.
3. Preliminaries
Contextual MDP We study contextual episodic MDPs, where each MDP Ma is associated with a context C E C belongs to the context space C. Further- more, Me. = {Ma,t consists of H different in- dividual MDPs, where each individual MDP Ma, (S,A, Pe,t (s'|s, Fe, (s)) Here S denotes the state space, A denotes the action space, Par denotes the transi- tion function and Fe.h denotes the reward function at stage h. We assume the starting state for each MM. is the same state T1. In this work, we interchangeablely use "environment" or MDP to denote the MDP Mr with different contexts.
Policy and value function We denote the policy Ty at stage h as a mapping S + (AA), which maps the current state to a distribution over the action space. We use т = [ /! h=
Ye et al. (2033) showed that in online RL, for a certain family of contextual MDPs, it is inherently im- possible to determine an optimal policy for each individual MDP. Given that offline RI. poses greater challenges than its online counterpart, this impossibility extends to finding optimal policies for each MDP in zero-shot offline RL setting well, which justifies our oplimization objective on the 7SG gap. Moreover, Ye et al. (2023) showed that the few-shot RL is able to find the optimal policy for individual MDPs. Clearly, such a setting is stronger than ours, and the additional interactions are often hard to be satisfied in real-world practice. We leave the study of such a setting for future work.
Remark 2
to denote their collection. Then for any episodic MDP M, we define the value function for some policy
(x)) == Eira +... + TH 1al x. an TW, The ~ Tw (sw, ax:), (sww,ax)))))))
8w ~ (- ax)), h1' 2 h),
(2, a) = Elra + ... + 12218h
ah a, Th Th (st, an) Aw Pv. - O2 -1;; ar TM,
Th Th (Sw, aw h 2 h- 1].
For any individual MDP M with reward r and transition dynamic P. we denote its Bellman operator [B a. fl (z, a) as
[B a) = B[rx(s (, a) + [(s')s'' ~ P(-Is,a)) Then
we have the well-known Bellman equation
VM)(()
Offline RL data collection process The data collection process is as follows. An experimenter i.idd samples num- ber n of contextual episodic MDP M, from the context set (e.g., ~ C). For each episodic MDP M,, the experimenter collects dataset D; = {<TT, oT rT,,, ) 1 1 which in- cludes K trajectories. Note that the action ai,,, selected by the experimenter can be arbitrary, and it does not need to follow specific behavior policy (Jin et al. 2021). We assume that D. is compliant with the episodic MDP Mi, which is defined as follows.
(QMA(E,)) (. (-)=)) Quaa(e,)) a) = [BM, Vi.++i(x,
For simplicity, we use Vihh Q Be,r to denote V VM Q. BMg.. We also use P. to denote P.M.. the joint distribution of any potential objects under the M. episodic MDP. We would like to find the near- optimal policy x² w.r.t. expectation, i.e., x² argmax, een Ve (I.),, where the set of collection of Markovian policies, and with little abuse of notation, we use Ee-- to denote the expectation taken w.r.t. the i.i.d.
3

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Definition 3 (Jin et al. 202T)) For dataset Dj :=
In general, H" can be very different from n² when the tran- sition probability functions of each environment are dif- ferent. For example, consider the 2-context cMDP prob- lem shown in Figure |I each context consists of one state and three possible actions. The offline dataset distributions I are marked on the arrows that both of the distributions are following near-optimal policy. By Proposition + in average MDP M the reward of the middle action de- terministically O, while both upper and lower actions deterministically 1. As result, the optimal policy #* will only have positive probabilities toward upper and lower ac- tions. This leads to E [V2,1 (21)] = O, though we can see that x is deterministically choosing the middle action and -o-c = 0.5. This theoretically illustrates that the generalization ability of offline RL algorithms without leveraging context information is weak. In sharp contrast, imitation learning such as behavior cloning (BC) converges to the teacher policy that is independent of the specific MDP. Therefore, offline RL methods such as CQL (Kumar et al., 2020) might enjoy worse generalization performance com- pared with BC, which aligns with the observation made by Mediratta al. 2023).
(o hh i. hihhh= h=1l==1 let DD be the joint distribution of the data collecting process. We say Dj is compliant with episodic MDP M, for any x' E &.r',T E [K] E [H], we have
PD. fr =T', - (
[(r) (
= =P.(rea(5h, on) =r', 8h+1 10 3h/ 0/ o.))
In general, we claim D; is compliant with M, when the conditional distribution of any tuple of reward and next state in Dr follows the conditional distribution determined by MDP M,,
4. Offline RL without context indicator information
In this section, we show that directly applying existing of- fline RL algorithms over datasets from multiple environ- ments without maintaining their identity information cannot yield a sufficient ZSG property, which is aligned with the existing observation of the poor generalization performance of offline RL. (Mediratta et al. 2023
polas) =1-6 ro(a1) = =1
(.((0))
21
rr(a)) =
In detail, given contextual MDPs MI,..., Mn and their corresponding offline datasets D1, .... Dni assume the agent only has the access to the offline dataset D = U D; where D = {(=x, A o,, A re h = ) Here C, E C' is the context information of trajectory T. which is unknown to the agent. To explain why offline RI. without knowing context information performs worse, we have the following proposition suggesting the offline dataset from multiple MDPs is not distinguishable from an "average MDP" if the offline dataset does not contain context information.
(a(as) ==
re(as)
=0 [a,(a))
[a
Th [a((2)) = 0
21
(22)
Bv Bu(as) =z rh (ag)
Figure I. Two Contextual MDPs with the same compliant average MDPs. The discrete contextual space is defined as C [u. w) and both MDPs satisfies S {w1 ).A {ai, a2, ag}, H 1. The data collection distributions and rewards for each action of each context are specified the graph.
Proposition 4 D is compliant with average MDP M (M MA = AS H, Fn),
Pa Jx, a a)Me. E Ba-ePo, (x,, a)
(x I., a) = Eaxe
5. Provable offline RL with zero-shot generalization
P(Fc.h = al).a P(fh = r|x, a)= == Eew Ee-CHo,(I
In this section, we propose offline RL with small ZSG gaps. We show that two popular offline RI. approaches, model- based RL and policy optimizaatonn-aased RL, can output RL agent with ZSG ability, with a pessimism-style modifica- tion that encourages the agent to follow the offline dataset pattern.
where ()) is the data collection distribution of (s,a) at stage h in dataset De-
Proof See Appendix A.1
Proposition 4 suggests that if no context information is revealed, then the merged offline dataset D is equivalent to a dataset collected from the average MDP M. Therefore, for any offline RL which outputs a Markovian policy, it converges to the optimal policy "* of the average MDP M.
5.1. Pessimistic policy evaluation
We consider a meta-algorithm to evaluate any policy л given an offline dataset, which serves as a key component in our
4

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Algorithm 1 Pessimistic Policy Evaluation (PPE)
Algorithm 2 Pessimistic Empirical Risk Minimization (PERM)
Require: Offline dataset (Da,)k=n Policy (TT)) confi-
dence probability 5 (0, 1).
Require: Offline dataset D
1: Initialize Vi (-) + 0, Vi E [n).
policy class II. confidence
2: for step h = H, H 1, 1,..., 1 do
probability 6 E (0, 1), a pessimistic offline policy evaluation
3: Let (B.,a (,
+ ODD,,, 1)
algorithm Evaluation as a subroutine.
1: Set De, = (<TT,, ol Ti,o +[(+))]
4: Set (..) + min - h + 1, (B.n ((+1))())))
2: π PERM = argmax, TEn R Di= (m)).
(1,a(,)))
5: Set (.) + (:-) m(+1))
where
( Evaluation (D,, ) x, 8/ (HHH((m)-)))
6: end for
7: Return D (...,..... Vin( (,,0),2
3: Return PERM
proposed offline RL with ZSG. To begin with, we consider general individual MDP and an oracle O which returns us an empirical Bellman operator and uncertainty quantifier, defined as follows.
compare with, we use the pessimism principle in the policy evaluation problem.
Remark 8 In our framework, pessimism can indeed facili- lale generalization, rather than hinder it. Specifically, we employ pessimism to construct reliable function for each environment individually. This approach supports broader generalization by maintaining multiple Q-networks sepa- rately. By doing so, we ensure that each function is robust within its specific environment, while the collective set of Q functions enables the system to generalize across different environments.
Definition 5 (Jin et al. 2021) For any individual MDP M, S 0. 1] that is compliant with M,
a dataset D S 8 test function VD C [0, HJS and confidence level &. we have an oracle ODD, VD that returns (BV
a tuple of Empirical Bellman operator and uncertainty quan- tifier. satisfying
Po (((VV))) - -BBMV))(r,
S {T(z,)) all (a, E x 2 1
5.2. Model-based approach: pessimistic empirical risk minimization
Remark Here we adapt a test function VD that can de- pend on the dataset D itself. Therefore, I is a function that depends on both the datasel and the test function class. We do not specify the test function class in this definition, and we will discuss its specific realization in Section 6
Given PPE, we propose algorithms that have the 7SG abil- ity. We first propose pessimistic empirical risk minimiza- tion (PERM) method which is model-based and conceptu- ally simple. The algorithm details are in Algorithm 2 In detail, for each dataset D; drawn from i-th environments, PERM builds a model using PPE to evaluate the policy x under the environment M, Then PERM outputs a policy π PERM II that maximizes the average pessimistic value. i..., 1/n Σ V. (21). Our approach is inspired by the classical empirical risk minimization approach adopted in supervised learning, and the Optimistic Model-based ERM proposed in Ye et al. (023) for online RL. Our setting is more challenging than the previous ones due to the RL setting and the offline setting, where the interaction be- tween the agent and the environment is completely disal- lowed. Therefore, unlike Ye et al. (2023), which adopted optimism style estimation to the policy value, adopt pessimism-style estimation to fight the distribution shift issue in the offline setting.
Remark 7 For general non-linear MDPs, one may employ the bootstrapping technique estimate uncertainty, in line with the bootstrapped DQN approach developed by (Osband et al. 2016). We note that when the bootstrapping method is straightforward to implement, the assumption of having access to an uncertainty quantifier is reasonable.
Based on the oracle O, we propose our pessimistic policy evaluation (PPE) algorithm as Algorithm D In general, PPE takes given policy as its input, and its goal is to evaluate the V value and Q value {(V.T, 2, ) of π on MDP M, Since the agent is not allowed to interact with M,, PPE evaluates the value based on the offline dataset { h= At each stage h. PPE utilizes the oracle O and obtains the empirical Bellman operator based on Di,, as well as its uncertainty quantifier, with high probability. Then PPE applies the pessimism principle to build the estimation of the Q function based on the empirical Bellman operator and the uncertainty quantifier. Such a principle has been widely studied and used in offline policy optimization, such as pessimistic value iteration (PEVI) (Jin et al., (02T). To
Next we propose a theoretical analysis of PERM. Denote AAT as the c-covering number of the policy space w.r.t. distance d(z1, x²) maxxes,keiil [- - 7 (-Is)]li. Then have the following theorem to provide upper bound of the suboptimality gap of the output policy л PERM
Theorem 9 Set the Evaluation subroutine Algorithm
5

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
as PPE (AlgoT). Let Li be the uncertainty quantifier returned by O Through the PERM. Then w.p. at least 6, the output x PERM of Algorithm 2 satisfies
Algorithm 3 Pessimistic Proximal Policy Optimzation (PPPO)
Require: Offline dataset
(D.
[< r(A) 1 confidence probability 8 E (0. 1),
2 2log(E (Hm) 1/5)
SubOpt(T π
PHRM
pessimistic offiine policy evaluation algorithm Evaluation as subroutine.
n
1: Set De, = { aith a, -- Ht)) [
Fish Lh+1 1+==
:Supervised learning eroo crror
2: Set To. as uniform distribution over A and 2o (.:)) as zero functions.
.. (T an)lsi = m)),
(1)
3: for = 1,2, ; do
1=
4: Set (-.)((1-)) x Ti T-1 exp(a Q211(())
Re:eeinforeement learning (RL) error
5: Set I......, N.....,,11 Q.1(1 (..... (-3(
where ,x is w.r.t. the trajectory induced by It with the transition Pe in the undertying MDP M,.
Evaluation( ((Dimattio(((,)))aa)))) nc, 5,, (/(nl)
6: end for
7: Return ा. PPPO = raaddo (TI,... nT)
Proof See Appendix
B.1
PERM in Algorithm 2 works as general model-based algo- rithm framework to enable ZSG for any pessimistic policy evaluation oracle. However, note that in order to implement PERM, one needs to maintain n different models or critic functions simultaneously in order to evaluate Σi= VA (21) for any candidate policy и. Note that existing online RL (Ghosh et al. 2021 achieves ZSG by model-free approach, which only maintains n policies rather than models/critic functions. Therefore, one natural question is whether we can design a model-free offline RL algorithm also with access only to policies.
Remark 10 The covering number (Hu)-1 depends on the policy class II. Without any specific assumptions, the policy class II that consists of all the policies x = {TT) ) hi-1 Th
N
S + ((A) and the log e-covering number log.? (((A))S) log(1 + {Alle))).
Remark 11 The SL error can be easily improved a distribution-depeddenn bound log.N. Varl where N is the covering number term denoted in 11 Var max, V Vo-VV2,(22 is the variance of the context distri- bution, by using a Bernstein-type concentration inequality in our proof. Therefore, for the singleton environment case where |C| 1, our suboptimality gap reduces to the one of PEVI Jin et al. )
We propose the pessimistic proximal policy optimization (PPPO) in Algorithm B address this issue. Our algorithm is inspired by the optimistic PPO (Cai et al. 2020) originally proposed for online RL. PPPO also adapts PPE as its sub- routine to evaluate any given policy pessimistically. Unlike PERM, PPPO only maintains n policies म Ti.... Tu, each of them is associated with an MDP Mr from the offline dataset. In detail, PPPO assigns an order for MDPs in the offline dataset and names them Mi,..., Ma. For i-th MDP M₁, PPPO selects the i-th policy π; the solution of the proximal policy optimization starting from T;-1, which is
Remark 12 In real-world settings, as the number of sam- pled contexts n may be very large, it is unrealistic to manage n models simultaneously in the implementation of PERM algorithm, thus we provide the suboptimality bound in line when when the offline dataset is merged into
with Theorem contexts such that < See Theorem 28 in A
Theorem 9 shows that the ZSG gap of PERM bounded by two terms In and I2. I1, which we call supervised learning error, depends on the number of environments n in the offline dataset D and the covering number of the function (policy) class, which is similar to the generalization error in supervised learning. I2, which we call it reinforcement learning error, is decided by the optimal policy x" that achieves the best zero-shot generalization performance and the uncertainty quantifier a In general, 12 is the "intrinsic uncertainty" denoted by Jin et al. (202I) over n MDPs, which characterizes how well each dataset D; covers the optimal policy n².
T.. f- argmax V (1111
11 KL [T|[T= = x1], (2)
where a is the step size parameter. Since ViT (21) is not achievable, we use a linear approximation Lę-1 (TT) to replace V (x1)) where
(.- = Vi + E. 186-1
H
( ) (u(lee)) o m- =x
(3)
5.3. Model-free approach: pessimistic proximal policy optimization
where QTHn x Q are the Q values evaluated on the offline dataset for M (2) and (3) give us a close-form
6

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
solution of π in Line 4 in Algorithm Such a routine corresponds to one iteration of PPO (Schulman et al. 2017)- Finally, PPPO outputs PPPPO as a random selection from ".... Tt-
measures Pi,h = (14 ,igh) over S and an unknown
vector 0e,h E IR. such that
Pa,n(a' (x,a) 10122, a), (i,h(z')),
(sr 8h 2, an a) = (0(2,0), Oc.)) (4) for all (x, E 8 x A x S at every step h E [H]
Remark 13 In Algorithum [ we adopt a data-sp data-splitting trick Jin et al.) 2021) to build Di, where we only utilize each trajectory once for one data tuple at some stage h. It is only used to avoid the statistical dependency of 1,h+1 (-) and 1 for the purpose of theoretical analysis.
We assume || a)ll S 1 for all (2,a) E S x A and
10 S vd at each step h E [H], and
we define | (S) = Is [1 dr.
We first specialize the general PPE algorithm (Algo to obtain the PPE algorithm tailored for linear MDPs (Algo This specialization is achieved by constructing B,,, V +
The following theorem bounds the suboptimality of PPPO.
Theorem 14 Ser the Evaluation subroutine in Algorithmm 3 as Algorithm ! Let sa be the uncertainty quantifier returned by O through the PPPO. Selecting a 1/V///HH Then selecting ō 1/8, wap. at least 2/3, have
T and ViJ based the dataset D;- We denote the set of trajectory indexes in Dr, as B, Algo subse- quently functions as the policy evaluation subroutine in Algo 2 and Algo for linear MDPs. In detail, we con- struct B., Ve.++1 (which is the estimation of By,,, V,,442 as (B.,, V VI, a) p(x,a) Wiih. where
log A
SubOpt( "Poy s 10
n
fy:SL everr
0i, = argmin EreB.. (TT + VAN+1(E7,))) - ó(II h. ). T "=)² + A- 1
B (sl =
(5)
1a 12: vrror
with À > 0 being the regularization parameter. The closed- form solution to (5) is in Line 4 in Algorithm 4 Besides, we construct the uncertainty quantifier 1s based on D, as
where Bi,x- is w.r.t. the trajectory induced by x²" with the transition P: in the underlying MDP M,
I,, a) = 3(0) 11o(2, A)llA,,,AA.
Proof See Appendix B.2
Σ
LW T I
Remark 15 As in Remark 12 we also provide the subop- timality hound in line with Theorem /44 when the offline dataset is merged into m contexts such that m < n. See Theorem 29 in C
r
with B(6) 0 being the scaling parameter.
Algorithm 4 Pessimistic Policy Evaluation (PPE): Linear
Theorem 14 shows that the suboptimality gap of PPPO can also be bounded by the SL error L and RL error Iz. Interestingly, / in Theorem 14 for PPPO only depends on the cardinality of the action space |A], which is different from the covering number term in J1 for PERM. Such difference is due to the fact that PPPO outputs the final policy PPPO as random selection from n existing policies, while PERM outputs one policy π PERM, Whether these two guarantees can be unified into one remains an open question.
MDP
Require: Offline
dataset (D.,) Danh policy n, confidence
{ T.. =.. )
probability 8 E (0.
1: Initialize V V (-) + 0, vi [
2: for step h = H, H I,..,, do
3: Set Ac, + EreB,n a( oLNooeT a[w) TT +
4: Set Winh + A ( Mao ah,a) (rīn +
((
5: Set Pan(s)) :) + 3(5) (Ø(
6. Provable generalization for offline linear MDPs
6:
Set Qin(..) + -minf((;) -)7 W.s (;,,(,)),, h
1)
7:
Set V.T + (Q)()()))
In this section, we instantiate our Algo 2 and Algo 3 for general MDPs on specific MDP classes. We consider the linear MDPs defined as follows.
8:
end for
9: Return V7 ())))).... (7,,),..,V (). (T, [..((.))
Assumption 16 (Yang and Wang 2019; Jin et al. 2019) We assume Vi E C,M, is a linear MDP with a known feature map o S x A-- + Rd if there exist d unknown
The following theorem shows the suboptimality gaps for Algo 2 (utilizing subroutine Algo and Algo (also with subroutine Algo
7

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Theorem 17 Under Assumption 16 in Algorithm 4 we set
Corollary 19 Suppose that for each E [n]. D, is gener- ated by behavior policy n, which well-explores MDP M, with constant Ci 2 Cmin- In Algo we set X = 1,, 1,6( d dH /log(AdH
A 1, B(8)=
= c- d.H v /log(2dH K
/O), where e > i a
positive constant. Then, we have:
(i) for the output policy n PERM of Algo with subroutine
where c' > 0 is a positive con- stant. Suppose we have K Z 40d/ Cmin log(Adn
A w.p. at least 8, the suboptimality gap satisfies
and
set C² := =1/n Ei= C₁ -1/2 Then we have:
20
anHNi 5
7 Tlog(E 110)
for the output T. PERM of Algo with subroutine Algo 4 w.p. at least 1 - 5, the suboptimality gap satisfies
SubOpt( )FRM) STV
(Bm)-
n
n
F [Io((nnaa)lls 81 = =z1)
(6)
2 2log(6A (Hm) (Mm) 1/5)
SubOpt( PERM STV
n
ffr the output policy N pPPO of Algo with subroutine AlgolA setting 5 = 1/88 then with probability at least 2/3, the suboptimality gop satisfies
+ 22726 8 H2 1/1
(og(12HHnKA((.. (Ba) 1/6) Cū, (8)
for the ouiput policy TPPO of Algo with subroutine Algald setting = -1/8, then with probability at least 2/3, the suboptimality gap satisfies
SubOpt(s ) & 10
( log LA|H2
B(
+
4mH
n
n
Es,or [Ili(sn; and an)l
$1 = ml) ):
SubOpt(7 (P0) S 10 log A
n
+2v26' 2 12.5 K 1-2 v /og(16dHnKK) ...HHnKk........ ...). (9) d
where Bi is with respect to the trajectory induced by n² with the transition Pi in the underlying MDP M, given the fixed matrix Ae, or Ns.
Remark 20 The mixed coverage parameter C₁ 1 E Vo. is large, ie.. the minimum eigenvalue of Σip ER,,M [0(N,, (nss, a is large. Amin(Es,,, indicates how well the behavior policy ã, ex- plores the state-action pairs on MDP M;; this shows that if for each environment i E [ni], the behavior policy explores Mi well, the suboptimality gap will be small.
1 is small if for any i E [n]. Gi
[0(S,, an) an))la indicates how well the state-action pair (sh) a1)) is covered by the dataset D:. The term E.s E Be,-- [1o(s,, an)]l 81 = =1 in the sub- optimality gap in Theorem is small if for each context i E [n]. the dataset Dj well covers the trajectory induced by the optimal policy n² on the corresponding MDP M,
n
Note that
Well-explored behavior policy Next we consider a case where the dataset D consists of i.i.d. trajectories col- lecting from different environments. Suppose D con- sists of n independent datasets D and for each environment i, D; consists of K trajectories D; {<II, a., rf,,) > independently and identically in- duced by a fixed behavior policy म; in the linear MDP M. We have the following assumption on well-explored policy:
Remark 21 Under the same conditions of Corollary
19
392 log(eN (Bn) 118) and K
(i) If
2
2
32d & 11 1og(122HnK (Rm) 18)
max 40d log( ddnH ) cwin then w.p. least - 8, SubOpt( PERM) Se.
}.
(i) If n 2
400H²² LAg(
and K
40d log(16dllf), 32e it log( (og(184H maxt can w.p. at least 2/3, SubOpt(==PPPO) Se.
}. then
2
Definition 18 (Duan et al. 2020; Jin et al. 2021) For an behavior policy मे and an episodic linear MDP M with feature map é, we say मे well-explores M with constant e if there exists an absolute positive constant 0 such that
Corollary 19 suggests that both of our proposed algorithms enjoy the O 1122 + K- C.)) convergence rate to the optimal policy x" given a well-exploration data collection assumption, where C; is mixed coverage parameter over n environments defined in Corollary 19
vh E [H], Amm) 2 c/d, where Σn = Br.M [d(s,, an (o(n)ss)))))) "I.
7. Conclusion and Future Work
A well-explored policy guarantees that the obtained trajecto- ries is "uniform" enough to represent any policy and value function. The following corollary shows that with the above assumption, the suboptimality gaps of Algo (with subrou- tine Algo 4 and Algo (with subroutine Algo decay to 0 when n and K are large enough.
In this work, study the zero-shot generalization (ZSG) performance of offline reinforcement learning (RL). We propose two offline RL frameworks, pessimistic empirical risk minimization and pessimistic proximal policy optimiza- tion, and show that both of them can find the optimal policy with ZSG ability. We also show that such generalization
8

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
property does not hold for offline RL without knowing the context information of the environment, which demonstrates the necessity of our proposed new algorithms. Currently, our theorems and algorithm design depend on the i.i.d. as- sumption of the environment selection. How to relax such an assumption remains an interesting future direction.
Anthony Cassandra, Lestie Pack Kaelbling, and Michael L Littman. Acting optimally in partially observable stochas- tic domains. In Aaai, volume 94, pages 1023-1028, 1994.
Karl Cobbe, Oleg Klimov, Christopher Hesse, Tachoon Kim, and J. Schulman. Quantifying generalization in reinforcement learning. International Conference On Machine Learning, 2018.
Impact Statement
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement
This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
learning via slow reinforcement learning, arXiv preprint /0li://611.027999 2016.
Yaqi Duan, Zeyu Jia, and Mengdi Wang. Minimax-optinal off aff-policy evaluation with linear function approximation. In International Conference on Machine Learning, pages 2701-2709. PMLR, 2020.
References
Anurag Ajay, Ge Yang, Ofir Nachum, and Pulkit Agrawal. Understanding the generalization gap in visual reinforce- ment learning. 2021.
Andy Ehrenberg, Robert Kirk, Minqi Jiang, Edward Grefen- stette, and Tim Rocktäschel. A study of off- off-policy learning in environments with procedural content gen- eration. In ICLR Workshop on Agent Learning in Open- Endedness, 2022.
Joshua Albrecht, Abraham Fetterman, Bryden Fogelman, Ellie Kitanidis, Bartosz Wróblewski. Nicole Sco, Michael Rosenthal, Maksis Knutins, Zack Polizzi, James Simon, et al. Avalon: A benchmark for rl generalization using procedurally generated worlds. Advances in Neural In- formation Processing Systems,
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree- based batch mode reinforcement learning. Journal of Machine Learning Research, 6, 2005.
2022.
Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhihong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforce- ment learning. arXiv preprint arXiv:2202.11556, 2022.
Chelsea Finn, Picter Abbeel, and Sergey Levine. Model- agnostic meta-learning for fast adaptation of deep net- works. In International conference on machine learning, pages 1 1126-11 1 1135, PMLR, 2017.
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A survey of meta-reiinorrcement learning. arXiv preprint arXiv:2301.0802. 080288 2023.
Kevin Frans and Phillip Isola. Powderworld: A platform for understanding generalization via rich task distributions. arXiv preprint 2r 2022.
Kamyar Ghasemipour, Shixiang Shane Gu, and Ofir Nachum. Why so pessimistic? estimating uncertain- ties for offline rl through ensembles, and why their in- dependence matters. Advances in Neural Information Processing Systems, 35:18267-18281, 2022.
Emmanuel Bengio, Joelle Pineau, and Doina Precup. Inter- ference and generalization in temporal difference learning. International Conference On Machine Learning, 2020.
Martin Bertran, Natalia Martinez, Mariano Phielipp. and Guillermo Sapiro. Instance-based generalization in re- inforcement learning. Advances in Neural Information Processing Systems, 33:11333-11344 2020.
Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan P Adams, and Sergey Levine. Why generalization in is difficult: Epistemic pomdps and implicit partial observability. Advances in neural information processing systems, 34:25502-255 225115, 2021.
Avinandan Bose, Simon Shaolei Du, and Maryam Fazel. Offline multi-task transfer r1 with representational penal- ization. arXiv preprint arXiv:2402.122577 2024.
Jiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li, and Liwei Wang. Near-optimal representation learning for linear bandits and linear rl. In International Conference on Machine Learning, pages 4349-4358, PMLR, 2021.
Emma Brunskill and Lihong Li. Sample complexity of multi-task reinforcement learning. arXiv preprint arXiv:1309.664 2013.
Qi Cai, Zhuoran Yang, Chi Jin, and Zhaoran Wang. Provably efficient exploration in policy optimization. In Interna- tional Conference on Machine Learning, pages 1283- 1294. PMLR, 2020.
Haque Ishfaq, Thanh Nguyen-Tane Songtao Feng, Raman Arora, Mengdi Wang, Ming Yin, and Doina Precup. Of- fline multitask representation learning for reinforcement learning. arXiv preprint arXiv:2403.1 1574, 2024.

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Scott Jeen, Tom Bewley, and Jonathan M Cullen. Conser- vative world models. arXiv preprint arXiv:2309.15178 2023.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offine reinforcement learning: Tutorial, review, and perspectives open problems. arXiv preprint arXi::2005.01643, 2020.
Yiding Jiang. Zico Kolter, and Roberta Raileanu. Uncertainty-driven exploration for generalization in re- inforcement learning. In Deep Reinforcement Learning Workshop NeurlPS 2022.
Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch off-policy reinforcement learning without great exploration. Advances in neural information processing systems, 33: 53:1264-1274, 2020.
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. arXiv preprint arXiv:1907.05388, 2019.
Rui Lu, Gao Huang, and Simon S Du. On the power of multitask representation learning in linear mdp. arXiv preprint S.XXv:2106.0805:: 2021.
Ying Jin, Zhuoran Yang. and Zhaoran Wang. pessimism provably efficient for offline ri? In International Con- ference on Machine Learning, pages 5084-5096. PMLR, 2021.
Clare Lyle, Mark Rowland, Will Dabney, Marta Kwiatkowska, and Yarin Gal. Learning dynamics and generalization in deep reinforcement learning. In Inter- national Conference Machine Learning, pages 14560- 14581. PMLR, 2022.
Arthur Juliani, Ahmed Khalifa, Vincent-Pierre Berges, Jonathan Harper, Ervin Teng, Hunter Henry, Adam Crespi, Julian Togelius, and Danny Lange. Obstacle Tower: A Generalization Challenge in Vision, Control, and Planning. In IJCAI, 2019.
Marlos C. Machado, Marc G. Bellemare, Erik Talvitie, Joel Veness, Matthew J. Hausknecht, and Michael H. Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. IJCAI, 2018.
Niels Justesen, Ruben Rodriguez Torrado, Philip Bontrager, Ahmed Khalifa, Julian Togelius, and Sebastian Risi. Illu- minating generalization in deep reinforcement learning through procedural level generation. arXiv: Learning, 2018.
Dhruv Malik, Yuanzhi Li, and Pradeep Ravikumar, When is generalizable reinforcement learning tractable? Advances in Neural Information Processing Systems, 34, 2021.
Bogdan Mazoure, Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Improving zero-shot generaliza- tion in offline reinforcement learning using generalized similarity functions. Advances in Neural Information Processing Systems, 0225088-22100 2022.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76:201-264, 2023.
Ishita Mediratta, Qingfei You, Minqi Jiang, and Roberta Raileanu. The generalization gap in offline reinforcement learning. arXiv preprint arXiv:2312.05742, 2023.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179-1191, 2020.
Thanh Nguyen-Tane and Raman Arora. On sample-efficin offline reinforcement learning: Data diversity, posterior sampling and beyond. Advances in Neural Information Processing Systems, 36, 2024.
Heinrich Küttler, Nantas Nardelli, Alexander H. Miller, Roberta Raileanu, Marco Selvatici, Edward Grefenstette, and Tim Rocktäschel. The NetHack Learning Environ- ment. In Proceedings of the Conference Neural Infor- mation Processing Systems (NeurIPS), 2020.
Alex Nichol, V. Pfau, Christopher Hesse, O. Klimov, and John Schulman. Gotta learn fast: A new benchmark for generalization rl. ArXiv, abs/1804.03720, 2018.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State- of-the-art, pages 73-73. Springer, 2012.
Ian Osband, Charles Blundell, Alexander Pritzel, and Ben- jamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Net- work randomization: A simple technique for generaliza- tion deep reinforcement learning. In International Con- ference on Learning Representations. https://openreview. net/forum, 2020.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähen- bühl, Vladlen Koltun, and Dawn Song. Assessing gener- alization in deep reinforcement learning. ICLR, 2019.
10

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Seohong Park, Tobias Kreiman, and Sergey Levine. Founda- tion policies with hilbert representations. arXiv preprint arXiv:2402.15567, 2024.
Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does zero-shot reinforcement learning exist? In ICLR, 2023.
Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. arXiv preprint arXiv:2107.062266 2021.
Lerrll Pinto, James Davidson, Rahul Sukthankar, and Ab- hinav Gupta. Robust adversarial reinforcement learning. In International conference on machine learning. pages 2817-2826. PMLR, 2017.
Masatoshi Uehara, Xuezhou Zhang, and Wen Sun. Repre- sentation learning for online and offline rl in low-rank mdps. arXiv preprint arXiv:1110004652, 2021.
Aravind Rajeswaran, Kendall Lowrey, Emanuel Todorov, and Sham M. Kakade. Towards generalization and sim- plicity in continuous control. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, Decem- ber 4-9, 2017, Long Beach, CA, USA, pages 6550-6561, 2017.
Huan Wang, Stephan Zheng, Caiming Xiong, and Richard Socher. On the generalization gap in reparameterizable reinforcement learning. In International Conference on Machine Learning, pages 6648-6658. PMLR, 2019.
Zhiyong Wang. Dongruo Zhou, John Lui, and Wen Sun. Model-based rl as a minimalist approach to horizon- and second-order bounds. arXiv preprint WrXiv:2408.0888 2024.
Paria Rashidinejad, Banghua Zhu, Cong Ma, Jiantao Jiao, and Stuart Russell. Bridging offline reinforcement learn- ing and imitation learning: A tale of pessimism. Ad- vances in Neural Information Processing Systems, 34: 11702-11716, 2021.
Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua Susskind, Jian Zhang. Ruslan Salakhutdinov, and Hanlin Goh. Un- certainty weighted actor-critic for offline reinforcement learning. arXiv preprint arXiv:2105.08140, 2021.
Shideh Rezaeifar, Robert Dadashi, Nino Vieillard, Léonard Hussenot, Olivier Bachem, Olivier Pietquin, and Matthieu Geist. Offline reinforcement learning as anti- exploration. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8106 -1668,44, 2022.
Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-connistent pessimism for offline reinforcement learning. Advances in neural infor- mation processing systems, 34:6683-6694, 2021a.
Martin Riedmiller. Neural fitted q iteration-firs experi- ences with a data efficient neural reinforcement learning method. In Machine Learning: ECML 2005: 16th Eu- ropean Conference on Machine Learning, Porto, Portu- gal, October 3-7, 2005. Proceedings 16, pages 317-328. Springer, 2005.
Tengyang Xie, Nan Jiang, Huan Wang. Caiming Xiong, and Yu Bai. Policy finetuning: Bridging sample-efficient offline and online reinforcement learning, Advances in neural information processing systems, 34:27395 2021b.
Mikayel Samvelyan, Robert Kirk, Vitaly Kurin, Jack Parker- Holder, Minqi Jiang, Eric Hambro, Fabio Petroni, Hein- rich Küttler, Edward Grefenstette, and Tim Rocktäschel. Minihack the planet: sandbox for open- reinforce- ment learning research, arXiv preprint arXiv:2109.13202, 2021.
Yuling Yan, Gen Li, Yuxin Chen, and Jianqing Fan. The efficacy of pessimism in asynchronous q q-learning. IEEE Transactions on Information Theory, 2023.
Lin Yang and Mengdi Wang. Sample-optimal parametric q-learning using linearly additive features. In Interna- tional Conference Machine Learning, pages 6995- 7004, 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Rad- ford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1rXiv:1707.063: 2017.
Rui Yang, Lin Yong, Xiaoteng Ma, Hao Hu, Chongjie Zhang, and Tong Zhang. What is essential for unseen goal generalization of offline goal-conditioned r!? In International Conference on Machine Learning, pages 39543-39571 PMLR, 2023.
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in rein- forcement learning. In International Conference Learning Representations, 2020. URL https://
Denis Yarats, David Brandfonbrener, Hao Liu, Michael Laskin, Pieter Abbeel, Alessandro Lazaric, and Lerrel Pinto. Don't 't change the algorithm, change the data: Ex- ploratory data for offline reinforcement learning. arXiv preprint arXiv:2201. 13425, 2022.
Andrea Tirinzoni, Riccardo Poiani, and Marcello Restelli. Sequential transfer in reinforcement learning with a gen- erative model. In International Conference on Machine Learning, pages 1482. PMIR,, 2020.

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Chang Ye, Ahmed Khalifa, Philip Bontrager, and Julian
Togelius. Rotation, translation, and cropping for zero- shot generalization. In 2020 IEEE Conference on Games (CoG), pages 57-64. IEEE, 2020.
Haotian Ye, Xiaoyu Chen, Liwei Wang, and Simon Shaolei Du. On the power of pre-training for generalization in rl: provable benefits and hardness. In International Confer- ence on Machine Learning, pages 19770-39800. PMLR, 2023.
Ming Yin, Yaqi Duan, Mengdi Wang, and Yu-Xiang Wang- Near-optimal offline reinforcement learning with linear representation: Leveraging variance information with pessimism. arXiv preprint S2KK:::203.058000 2022.
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 re- inforcement learning. Advances in neural information 2021.
processing systems,
Amy Zhang, Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization in continuous reinforce-
ment learming. ArXiv, abs/1886.07997, 2018a
Chicheng Zhang and Zhi Wang. Provably efficient multi- task reinforcement learning with model transfer. Ad- vances in Neural Information Processing Systems, 34, 2021.
Chiyuan Zhang, Oriol Vinyals, Rémi Munos, and Samy Bengio. A study overfitting in deep reinforcement learning, ArXiv, abs/1804.06893, 2018b.
Weitong Zhang, Jiafan He, Dongruo Zhou, Amy Zhang, and Quanquan Gu. Provably efficient representation selection in low-rank markov decision processes: from online to of- fline rl. In Proceedings of the Thirty-Ninth Conference Uncertainty in Artificial Intelligence, pages 2488-228 2497, 2023.
12

---

Provable Zero-Shot Generalization in Offlin Reinforcement Learning
We provide missing proofs and theoretical results of paper in the Appendix sections:
we provide the missing results of Section 4 We first provide the proof of Proposition then we analyze the suboptimality gap of the Pessimistic Value Iteration (PEVI) (Gin et al. 2021)) in the contextual linear MDP setting without context information.
In Appendix A
In Appendix
B provide the proofs of our main theorems the suboptimality bounds of PERM and PPPO in Section
5
In Appendix state and prove the suboptimality bounds we promised in Remarks 12 and 15 where we merge the
sampled contexts into groups (m < n) reduce the computational complexity in practical settings.
In Appendix D we provide the proofs of results in Section on linear MDPs. Namely, we provide proof of Theorem 17
proof of Corollary T9
A. Results in Section 4
A.1. Proof of Proposition 4
Let = ((32 A. oe, Ah Te (h) )== denote the merged dataset, where each trajectory belongs to context C-. For simplicity, let D. denote the collection trajectories that belong to MDP Me. Then each trajectory D' is generated by the following steps:
The experimenter randomly samples an environment C.
The experimenter collect trajectory from the episodic MDP Mc.
Then for any x', ,,,,T we have
P =r' r' = ([ oz 11j=1 fod sy xd s4 )==tt
(r = 141 = x', 11 n N -=1 tr 2e
(D oy d }1 fr' Fre -t ا
= -E Po (re, = = h+l = x 2/11 og )) fre 25 c, }ز= o = c)g(c),
(10)
0EC
where
a(c) = (oo(( Re (oyn))s=1 Paghn (ل
N) Evec (D en Le N)} (re hi
1--1 )
Next, we further have
10
= = = =z'ssh = Fe,,,,
Σ
te. ap = (g(c)
cEC
(s)) = = = x11sh = xx- x: = a a=,,,, n (8h = x h, ah = σ nr c)
EcEr PD (8h = w an = a_ 22 =
eEC
Pa(ran(Hn) = =r', 8h41 884 = 1'13h = xt T2 an = a x )P = x 82, no
= of ar
pp(c) p(c)
E.so p(c) Po(sh = x ah T.-
cEC
EE-o Polra,((s)) =
, 8h+11 1'18h xo F Oh a a Wile ( (2or Ep-CHehh (x a²
where the first equality holds since for all trajectories T satisfying C, = c, they are compliant with Mey the second one holds since trajectories are independent of each other, the third and fourth ones hold due the definition of (e,a())))
13

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
A.2. PEVI algorithm
Algorithm 5 (Jin et al. 2021) Pessimistic Value Iteration (PEVI)
Require: Dataset D [(02 4 o Ae Te Whit confidence probability 5 E (0.1).
1: Initialization: Set (+)+ +- O.
2: for step = H. ,,..,,ldo do
3: Set An -ー= oi5 (7) ( 1..I.
4: Set Φh t - o( [T) (r) + V1+1(11411))
Set (n(:;-) < -((() ((5). ()(:)) TA, )
5:
6:
Set (x(:)) t min(((;.) in -Tass H -
-h+ h 13*. () + argmax,,, vn. (Qn(;,); (-1 (.)))))
7:
Set (-)-))
8:
Set (-) - (.)()))) ( (.1.)).
9:
end for
10: Return PEVI = 1
We analyze the suboptimality gap of the Pessimistic Value Iteration (PEVI) (Jin et al. 2021 in the contextual linear MDP setting without context information to demonstrate that by finding the optimal policy for M is not enough to find the policy that performs well on MDPs with context information.
Pessimistic Value Iteration (PEVI). Let Ā² be the optimal policy w.r.t. the average MDP M. We analyze the performance of the Pessimistic Value Iteration (PEVI) (Tin al. 202T under the unknown context information setting. The details of PEVI is in 5
Suppose that D consists of K number of trajectories generated i.i.d. following by a fixed behavior policy म. Then the following theorem shows the suboptimality gap for 5 not converge to 0 even when the data size grows to infinity.
Theorem 22 Assume that ãt In Algo we set
X = B(6) = d' dH (og(AAHK/)))
(11)
where è' > O is positive constant. Suppose have K ē. d logg( (AdH/E), where ō > 0 is sufficiently large positive constant that depends on Then we have; w.p. at least - for the oulput policy PEV of Algo
sup VM
VM
PSW S d" 8 22222KK 1/2 /oog(AdHKK)))
(12)
and the suboptimality gap satisfies
SubO d" 6//2 (og(AdHK//)) + 2sup VM, (x1) Eewo Vo (x:),,
(13)
where c" > is positive constant that only depends on e. and d.
Proof [Proof of Theorem 22 First, we define the value function on the average MDP M as follows.
H
VT(()) Er.M
Σ
re(ss, az) sh = x
(14)
We then decompose the suboptimality gap as follows.
SubOpt(77PEV)
Eavo [ (x1) Eenc V.a (21)
PEVI
=v (21) v non
(x1) + (Eowc [VT] (z1) - V. (21) + (Vi (x1) -eewo [VC]] mii
(x1)))
< VV (x1) - Vi u
(x1) + 2sup IR (x1) Eoro V2l (21)).
(15)
14

---

Provable Zero-Shot Generalization in Offine Reinforcement Learning
Then, applying Corollary 4.6 in (Jin et al.
2021), can get that w.p. least -
TMvI
Vi
(11)
(.1) S d' /3/2
1/2 1///////oog// (og(4HHH//))
(16)
which, together with Eq. completes the proof.
Theorem 22 shows that by adapting the standard pessimistic offline RI. algorithm over the offline dataset without context information, the learned policy π PEVI converges to the optimal policy म" over the average MDP M..
B. Proof of Theorems in Section 5
B.1. Proof of Theorem 9
We define the model estimation error
ax, = (Bir, ( - Q (z.
(17)
And define the following condition
(B.,A (2) - (BB,, (ih+)) a) s (2, a) for all i e in], (2,a) 8x A, [H]
(18)
We introduce the following lemma to bound the model estimation error.
Lemma 23 (Model estimation error bound (Adapted from Lemma 5.1 in (Jin et al. 2021 Under the condition of Eq. 18 we have
S (2, s (x, for all i [n], म E 11, (2,a) E ex A, he
(19)
Then, prove the following lemma for pessimism in V values.
Lemma 24 (Pessimism for Estimated V Values) Under the condition of 18 for any i e [n], E II, E S, we have
(x) (x).
(20)
Proof For any i E [n], II, E S,a A, we have
(x, (x,
2 runl a) (T (z, (rin a) + (Br. Vr Vi - a))(z,a))
= (Br, Vi - (B., Dx 1) a) + +Tan(z,a)
(B BI, (z,
2 (B. (Z - 1,6+ + (z,
= (Po, - V(h+1) (2. a),
where the second inequality is because of Eq.18 And since the H step we have VH Va 0, we can get (z, a) (x, a). Then induction prove Q a) 2 (x, for all h. Given Q7, 2
18
(2,41144, u), we have
(In(x,a) ax) 2 (Prn(V.T (r,+t))
E ((Ln))(0)t1))) Q7 (sh+1, :) The
x, = (119441)) (
0.
(21)
15

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Then have
(x) (x)) = (a,, 27,6 ( 2 0.
Then start our proof.
Proof [Proof of Theorem
First, we decompose the suboptimality gap as follows
SubOpt(T PERM)
BecoV (x1) - c.1 (x1)
= BosoVV, (x1)
1.1
(x1)
+
Σ
1,1
PERM
(x1) - Eewe Ve,1
FERM
(x1)
1=1
1=1
( (x1) (x1)).
Vi,
PERM
(22)
n
For the first two terms, we can bound them following the standard generalization techniques Yee et al. 2023) ie., we use the covering argument, Chernoff bod union bound.
Define the distance between policies d( x²) - max,cs,Me[M [|x](1))))))]))))))){{)) 71(115)112 We construct the e-covering set II w.r.t. such that
VT IT, 3. E IT s.d. d(a, se
(23)
Then we have
Vi [n], 11, Fh E Ii st.V (x1) I He.
(24)
By the definition of the covering number, with prob. least 1 $. for any II.
Ī = =Ni By Chernoff bound and union bound over the policy set îI, we have
2 log log((6N
v (m)) Eono
Vá (01) SV
(25)
n
By Eq. 24) and Eq. 25).
Vi [n]. 11 E ñî with i NTT s.t.V (2)) - (2)(( s He, and with probability at
least 1 /3, we have
VA (x1) Eeno (21(
V (s1) E Vo
n
1 Σ v (81) Vi (51)
+ Eaw
BanoV V (2, - Eanc V2
log (NT188) 2HE
(26)
n
Therefore, we have for the first two terms, w.p. at least
25 we can upper bound them with 4He + 2/
2 2log
16

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Then, what remains is to bound the term 1. - (V (x1) (21)).
V. M
First, by similar arguments, we have
va
(21) - Vii
MRM
(x1) - (V, (x1) - V.21 vM (x1)) + 1V 1 MIRM
(x11 - V. PRRM
(r1)
< He + VIi (a1) V.Tin (x1):
(27)
where PERM E IT such that [V,1 (x1)
Vi
" PERM (x)
S He.
By the definition of the oracle Definition 5 the algorithm design Algo (e.g., we call oracle (D2, (Hx)) +))). and use union bound over H steps, contexts, and V (Hn) policies, we have:
with probability at least 6/3, the condition Eq.18) holds (with the policy class IT replaced by II (and e = 1/(H))).
Then, have
1,1 (x1) - mu (I1))
(V (x1) -Pri Dijii (2i))
n
(VTT (x1) -Vi (xi))
(x1) - -.. (21))
i=1
(VTT (x1) -PzT Piimeau (x1)) H
Hn
(i (x1) V (z1)) + 1/n,
(28)
where the first inequality holds because of the pessimism in Lemma 24 the second inequality holds because S V
(x1)
PERM
(21) < He with here specified (/H,,, and the last inequality holds because that in the algorithm design of we set x PFRM argmax 1 n Σm= V (21).
Algo
Then what left is to bound VI (xz) - V (x1).
And using Lemma A.I in (Jin et al. 2021), we have
H
H
V.5 (x1) - (x1) =
(xz)
Σ
Eq M (sn, = x| +
Σ
Ex M. .M (iTh(sh, = z])
Aml
h=1
1. E B MM 1 (-) (-88)) - m,( (a))a 81 x²]
4=1
H
< 2
Ex .M, (Tn, an) x]
MM
(29)
1=1
where in the last inequality we use Lemma
23
Finally, with Fq.22. 22). Eq. 26
Eq.27
27
Eq.
28
and Eq. 29
with e set as nH. we can get w.p. least 1
17

---

Provable Zero-Shot Generalization in Offlin Reinforcement Learning
Berc Vx (x1)
V2,1 (x1)
2log(6NV (Hn) 1/6)
H
5 +
E M. an) = z1]
1=1 Aml
2
1lgg(AA (11)
2 n
7V
B MM (en(on an)ls1 x1)
1=
B.2. Proof of Theorem
Our proof has two steps. First, we define that
a, B., (.A (2,
(30)
Then have the following lemma from Jin al. (2021)
Lemma 25 Define the event & as
DaL (1, s E ( (B,,,
[M), vi e [m]
Then by selecting the input parameter & 6/((n) in O, have IPE)) 1 8 8 and
0 S (2, < 21 (x.
Proof The proof is the same as [Lemma 5.1, Jin et al. 2021 with the probability assigned as 6/(H)) and union bound over h E i i E [n).
Next lemma shows the difference between the value of the optimal policy n' and number n of different policies n; for 7 MDPs.
Lemma 26 Let be an arbitrary policy. Then we have
#
Σ (m))
E. E. (.ni((u(s)) (-)))-)))al
i=
1=1
(E 14,, an)] Ei 14,n ax)))
(31)
Proof The proof is the same Lemma 3.1 in (Jin al. 2021) except substituting into the lemma.
We also have the following one-step lemma:
Lemma 27 (Lemma 3.3, Cai et al. 2020) For any distribution p",p A(A), ifp'(-) x p(.) exp(a Qax, -), then
(Q(x,)) p()) S aH2 (KL(O"())P() -MLL)(O)O)OKK-XXX))>))))))))))))
Given the above lemmas, we begin our proof of Theorem
18

---

we have Provable Zero-Shot Generalization in Offline Reinforcement Learning
Proof (Proof of Theorem
Combining Lemma 25 and Lemma 26 have
(x)) (
H
1 T.ta) +
Bi [Q `
Bi (T),
1=1
B [ (en( an]]
H Σ h=1
n
H
I SaH + a-
Bi
KL
+ +
Σ B (2h, ax))
i=1
11
S salf'n/2+0 11 log
|A| +
B ITa (vn, an)],
h=1
where the last inequality holds since 41,h is the uniform distribution over A. Then, selecting
> Iia (x1) - V (z1) < 2 log A +2 B TTn an)l,
1=1
which holds for the random selection of 'D with probability at least 8. Meanwhile, note that each MDP M; is drawn ii.d. from Meanwhile, note that n; only depends on MDP M1, M Therefore, by the standard online-to-batch conversion, we have
P ( ( -
(=))
Σ
n
Bew Vo, Ecwo Val (mi) - 2H
2 log 1/5
2 1 -
which suggests that with probability at least - 25,
n Bew (x1) s 2 log |A
2 n E (Tn an) +2
2H log 1/8
Eewo ViT (x1)
n
n
Therefore, by selecting x PPPO random (T....,, and applying the Markov inequality, setting 1/8, have bound holds.
C. Suboptimality bounds for real-world setups
In this section state and prove the suboptimality bounds we promised in Remarks 12 and 15 where merge the sampled contexts into groups (generally, << n) to reduce the computational complexity in practical settings. The bound in Theorem 28 serves as partial justification for the effectiveness of IQL-mV V in our real-data experiments (Section ??).
Assume m In and the n contexts from offline dataset are equally partitioned into m groups. We write the resulting average MDPs (see Proposition 4) for each group as M Mm. For each Mj. we regard it as an individual context in the sense of 18 and denote the resulting uncertainty quantifier and value function as I" j.h, V j.b.
Theorem 28 (Suboptimality bound for Remark 12)
Assume the same setting as Theorem 9 with the original n contexts grouped contexts, and denote the resulting agorithh PERM-mX V. Then w.p. at Teast ō, the output n' of PERM-mV V satisfies
19

---

Provable Zero-Shot Generalization in Offlie Reinforcement Learning
2 llgg (1mm)
H
SubOpt(aT) 2
B
M, [T' s,n an |s1 = æi)
n
A=
SSpperiized learning
Ie:Reinfoveement learning (RL) crror
+ m
+
2sup
1 (mi)
VTT
TI
m
Addininnl opprosimation
where Ey. is w.E.t the trajectory induced by x² with the transition P, in the underlying average MDP Mj.
Proof [Proof of Theorem Similar to the proof of Theorem
we decompose the suboptimality gap as follows
SubOpt(x')
Vi (21) EEacV V 0,1 (11)
E Vel (21)
n
V (x1)
+
Vi (x1) REmoVVil
n
m
Vil (I1)
Au (21)
Vi
1. (x1)
1 Σ
ViT((a) (21)
V
(
V (z1))
(32)
m
Note that we can bound the first and third lines of 32 with the exactly same arguments as the proof of Theorem 9 the only notation- wise difference is that the uncertainty quantifier becomes T' as we are operating on the level of average MDP M,,
The only thing left is to bound the second line of 332 This is the same in spirit of the bound [5p. so that we can express the bound follows
n
(x1)
1
m
V (x1)
1 m
(x1)
Σ n
VA (x1)
n
m
m
< suu
(x1)
3.1 (x1)
To conclude, our final bound can be expressed with mi can get w.p. at least ō
SubOpt(TT)
2 loggGNN1 (Hm) (Hm)-1/0) ./6)
H
2
21
+
E MM [T'
('nn(sn an)
1=1
1 n
m
+ 2 sup
V (x1)
Σ
V (21)
m
n
m
1=1
To prove the suboptimality bound for Remark 15 we denote that the policies produced by PPPO after merging dataset to m groups to be T....., 7 and the original PPPO algorithm would produce the policies as ,',..., T//- We assume that the merging dataset from n to m groups is only to combine the consecutive n m terms from , and preserves the order.
20

---

Provable Zero-Shot Generalization in Offline Reinforcement Learning
Assume the same setting Theorem 14 witth the original contexts grouped as m contexts, and denote the resulting algorithm as PPPO-mV. Let I j.h be the uncertainty quantifier returned by O through the PPPO-mV algorithm. Selecting a 1/v H... Then selecting 8 = 1/8, w.p. at least 2/3, we have
Theorem 29 (Suboptimality bound for Remark 15)
SubOpt(z=" (P Sub0pu(("700-/)) S 10
log |A
1
H 3 [T' (sh, mi
V
m
3=1 h=l
12:RL orny
+ sup
V. (11)
m Σ V (21)
1
+ Σ sup Ee Ee(V -V (xx))
71
sup [V c, (+1)] - V' (((() (Tn))
T1
where Ex is w.r.t the trajectory induced by x² with the transition P, in the underlying MDP Mj.
Proof [Proof of Theorem Using the same arguments as in the proof of Theorem 14 with a = 1/VH2m, we can derive the bound
m
H R5 3. an)].
Σ
(iVi - (
2/m mm |||H2 +
==1
j=1
Leveraging this bound and anline-to-batche we obtain the following estimation
1 m
E. (z1)]
Eg V V (x1)
m7
1 Σ
- Σ
1 m
-E 01 (x1))
Bel c.l
[x1)]
+ m
Re (
Eg V.li (I1)
m
i=1
j=1
$2H
2 2 1/6 1
Ecl E. (z1) ( - viJ (2z))
Σ 1=1
Vi (x1)
Σ j=1
E. (.
-2H
2 2oo 1/6
1
Vil (x1)
m Σ
1 m
v 1,1 (21)
n
1 m
1 m
VIT"
(,1(z)) (z1)
V (21)
m
m
m Ee [[]i)
Be IV (o (x) )] v()(()))
+ m
V (x1)
71
m j=1
i=1
j=1
2 log 1/6
1
<2H
sup
ViT
V
5.1 (x1)
π7
1=1
==1
log |
2
H
2/
E Ej.r T (sn,
m
m
1 n
1 m m
sup VV -v (z1)
L.lp/E.lV' e,1 (21) V (x1))
Finally we apply Markov inequality and take 6 = 1/8 as in the proof of Theorem 14
21

---

Provable Zero-Shot Generalization in Offlin Reinforcement Learning
D. Results in Section 6
D.1. Proof Theorem
By (Jin et al. 2021). the parameters specified as 1, {(6) = c. dH (/) and applying union bound, we can get: for Algo with probability least -8/3
[Bi, Vi. ( (2,, (Ber, V (z, a)l B
(o(z, ( ())
3nHA (Hm)-1
for all i [n]. iT E ii, (z. a) E E x A.K [
(33)
where II is the Hn -ooering of the policy space II w.r.t. distance d(r x²) max ||/ (-)(--)))l (-1s)ll1-
Therefore, we can specify the (:.) in Theorem with B JnHA (0(1, A,,,0 1/2 and follow the same
process as the proof of Theorem to get the result for Algo 2 with subroutine A
Similarly, we can get: we can get: for Algo with probability at least - 1/4
(B., (ant))( (Ba, (ittt)) (2, 3 4nH (0(x, ( 1/2 for all i E [m], a) S A, [H].
(34)
Therefore, we can specify the (:)) in Theorem 14 with B (a (Ø(z, a) T) A M0,, ()) and follow the same process as the proof of Theorem 14 to get the result for A with subroutine Algo
D.2. Proof of Corollary 19
By the assumption that D; is generated by behavior policy म, which well-exploree MDP M; with constant (where the well-explore is defined in I8) the proof of Corollary 4.6 in (Jin et al. (021), and applying a union bound over contexts, we have that for Algo with subroutine Algo w.pp at least 8/2
2d
( $ C,K
for all i [n], (I, a) E S x A and all h [4],
(35)
and for A with subroutine Algo w.p. at least 1 5/2
2dH
16(x,,))lA; $ for all i [n], (I, S A and all h E [H]
K.K
(36)
because the data splitting technique and only utilize each trajectory once for one data tuple at stage h, so replace K with K/H.
Then, the result follows by plugging the results above into Theoren 17
22


Reviews:
[{"review_id": "PjJcfYvs6J", "reviewer": "Reviewer_KVto", "summary": "Main findings&Contribution: \n1. Problem finding: Classic offline reinforcement learning methods often struggle to generalize unseen environments, primarily due to dataset coverage limitations and the absence of contextual information during training. Merging multi-environment datasets without preserving contextual distinctions may result in convergence to an \"average MDP,\" resulting in poor zero-shot generalization (ZSG).\n2. Main contribution: This paper proposes Pessimistic Empirical Risk Minimization (PERM) and Pessimistic Proximal Policy Optimization (PPPO). These methods leverage pessimistic evaluation(PPE) to balance data fitting and generalization.\n3. Theoretical Guarantees: The suboptimality gap of PERM/PPPO is jointly determined by the supervised learning (SL) error (governed by the number of training environments n and dataset size) and the reinforcement learning (RL) error (influenced by offline data coverage quality and uncertainty quantifiers).", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks very much for reviewing our work and for the valuable suggestions. Our responses are as follows.\n\n**Q1:**  Augment theory with experiments; necessity of pessimism\n\n**A1:** Thank you for the thoughtful suggestion. We would like to respectfully emphasize that this is a theoretical paper submitted to the theory category, and our primary goal is to develop first-principles theoretical insights into zero-shot generalization in offline RL—a significantly under-explored area compared to standard offline RL in a single environment or meta-RL approaches that rely on additional adaptation data. We believe that our theoretical contributions represent a substantial advancement in the foundations of reinforcement learning and are valuable in their own right. We plan to run a thorough experimental evaluation of our proposed PERM and PPPO algorithms in future work.\n\nFor the necessity of pessimism, we conducted a small‐scale experiment on a Combination Lock environment with 10 contexts adapted from (Bose et. al, 2024). We use 500 exploratory trajectories from (Agarwal et. al, 2023) per context to form the offline dataset.\n\nThe table reports the average returns of the learned policies (evaluated over 100 runs) for our PPPO algorithm (with and without pessimism). We see that standard PPPO achieves a higher average reward, validating our theoretical premise that using environment‐specific pessimism can improve zero‐shot generalization.\n| PPPO           | PPPO w/o Pessimism |\n|--------------------|------------------------|\n| 0.0650 ± 0.0173    | 0.0603 ± 0.008         |\n\n\nAlekh Agarwal, Yuda Song, Wen Sun, Kaiwen Wang, Mengdi Wang, and Xuezhou Zhang. Provable benefits of representational transfer in reinforcement learning. In The Thirty Sixth Annual Conference on Learning Theory, pages 2114–2187. PMLR, 2023.321\n\n\n\n\n\n**Q2:**  The paper argues that per-environment pessimism avoids conflating dynamics across contexts. This is intuitively justified but not empirically validated. The theory assumes environments are i.i.d., which is critical for the SL error term but not thoroughly motivated (e.g., no discussion of non-i.i.d. settings).\n\n**A2:**  Thank you for the insightful feedback. Our per-environment pessimism approach ensures that each environment’s offline data contributes to a separate (pessimistic) critic or model component. This preserves environment-specific information, which is leveraged for generalization—either by maximizing the averaged value (PERM) or conducting PPO-like updates across environments (PPPO). By focusing pessimism on individual environments, we avoid the pitfalls of simply averaging across distinct dynamics and rewards, which could undermine generalization, as we theoretically demonstrate in Section 4. Importantly, this approach is not only intuitive but also forms the basis for the theoretical guarantees presented in Theorems 9 and 14. Additionally, as mentioned in **A1**, we have conducted some empirical analysis to support these findings.\n\nRegarding the i.i.d. assumption, we acknowledge that generalization in offline RL under distribution shifts remains an open challenge, and we are happy to discuss more in revision. Prior work on RL generalization in the online setting (Ye et al., 2023) also assumes i.i.d. environments and highlights OOD generalization as an open question. We believe starting with the i.i.d. assumption—as done in prior work—is a reasonable first step, and we view extending this analysis to non-i.i.d. settings as an exciting future direction.\n\n\n**Q3:**  The RL error term in the suboptimality gap scales linearly with the state–action space dimension  $d$. Could this dimensional dependency result in significant performance degradation in practice? Have you considered incorporating representation learning (e.g., low-rank feature extraction) to mitigate the curse of dimensionality?\n\n**A3:**  We believe you are referring to Theorem 17, which analyzes the special case of linear MDPs. We respectfully clarify that the parameter $d$ here refers to the dimension of the linear feature map, as defined in Assumption 16, and **not** to the cardinality of the state space $|\\mathcal{S}|$ or action space $|\\mathcal{A}|$. In fact, this formulation inherently assumes a low-dimensional representation, where the true dynamics and rewards lie in a $d$-dimensional linear subspace. Without such an assumption, we agree with you that incorporating representation learning is necessary. Note that our work could incorporate with recent representation learning works (Ishfaq et al., 2024) and taking the low-dimensional learned representation as our features. \nWe will clarify this point further in our revision to avoid any confusion following your advice.\n\nAgain, thank you very much for reviewing our paper and for the thoughtful feedback. We hope our responses adequately address your concerns, and we would be happy to clarify further if needed.", "claims_and_evidence": "**1. Claim:** Standard offline RL fails at ZSG without context information because the dataset resembles an \"average MDP.\"\n\n**Evidence:** Proposition 4 theoretically shows that merging datasets across contexts creates an indistinguishable average MDP. A 2-context example (Figure 1) illustrates how optimal policies for individual environments differ from those for the average MDP.\n\n​​**2. Claim:** PERM and PPPO achieve ZSG with suboptimality bounds decomposed into SL and RL errors.\n\n**Evidence:** Theorems 9 and 14 bound the suboptimality gap for both algorithms. The bounds depend on policy class complexity (SL term) and dataset coverage (RL term). While the focus is theoretical, even simple experiments (e.g., tabular MDPs or linear function approximation) could validate the practical relevance of the bounds and the necessity of pessimism.\n\n**3. ​Claim:** Pessimism enables generalization by preventing overfitting to individual environments.\n\n​**Evidence:** The paper argues that per-environment pessimism avoids conflating dynamics across contexts. This is intuitively justified but not empirically validated. The theory assumes environments are i.i.d., which is critical for the SL error term but not thoroughly motivated (e.g., no discussion of non-i.i.d. settings).", "methods_and_evaluation_criteria": "The methods are theoretically grounded and address the core challenges of ZSG in offline RL. The evaluation criteria are appropriate for a theory-focused paper but would benefit from broader empirical validation to demonstrate practical utility.", "theoretical_claims": "The theoretical claims are ​largely correct and build thoughtfully on existing offline RL frameworks. The proofs leverage established techniques (pessimism, covering numbers, elliptical potentials) while extending them to the multi-environment ZSG setting. However, tighter characterization of environment diversity (beyond i.i.d. assumptions) and empirical validation would elevate the work from a theoretical contribution to a practical guide for ZSG in offline RL.", "experimental_designs_or_analyses": "This paper does not provide experimental validation for the proposed methods.", "relation_to_broader_scientific_literature": "By integrating pessimism with multi-environment policy optimization, the paper bridges gaps between offline RL, generalization theory, and contextual decision-making, offering a principled framework for ZSG with implications for real-world applications where environment-specific interaction is costly or impossible.", "essential_references_not_discussed": "The related work has been well-included in the main text of the paper.", "other_strengths_and_weaknesses": "While the paper provides a rigorous theoretical analysis of ZSG in offline RL, the empirical validation of the proposed methods (PERM and PPPO) is limited in scope. To strengthen the practical relevance and credibility of the claims, the authors should augment their evaluation with ​cross-task generalization experiments on widely adopted benchmarks like D4RL[1].\n\n[1] Fu, J., Kumar, A., Nachum, O., Tucker, G., & Levine, S. (2020). D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219.", "other_comments_or_suggestions": "No other comments and suggestions.", "questions_for_authors": "Q1. The reinforcement learning error term in the suboptimality gap scales linearly with the state-action space dimension $d$. For high-dimensional tasks, could this dimensional dependency result in significant performance degradation in practice? Have you considered incorporating representation learning (such as low-rank feature extraction) to mitigate the curse of dimensionality?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wr6RpQipIV", "reviewer": "Reviewer_qPEC", "summary": "The work theoretically analyzes contextual reinforcement learning in the offline setting for zero-shot generalizability. In particular, the work proposes two algorithmic frameworks with provable zero-shot generalization abilities.\n\n\n## update after rebuttal\nI have read all reviews and am keeping my score. I do believe that the work is interesting to a good chunk of the (RL sub)community and would lead to interesting discussions on acceptance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks very much for your positive feedback and your valuable suggestions. Our responses are as follows.\n\n**Q1:**  One point I remain uncertain about is the assumptions about the context at hand. The work seems to implicitly assume perfect context information. It would be best to highlight this assumption.\n\n**A1:**  Thank you for raising this important point. We respectfully clarify that our approach does not assume access to perfect or rich context information. Specifically, our algorithm and analysis only require that, during training, we can identify which trajectories come from the same environment. That is, trajectories from the same environment must be grouped together under a shared context label—this label can be **arbitrary** and need not encode any meaningful or descriptive features of the environment. It merely serves to distinguish between different environment instances during training.\n\nImportantly, **no context information is required at test time**. The agent is expected to generalize to a new environment without observing any context label or side information.\n\nAs shown in Section 4, without such grouping, the data effectively collapses into an average MDP, which leads to poor zero-shot generalization—a limitation of prior offline RL methods. We agree that this assumption could be more explicitly stated and will revise the paper to clarify it. Investigating whether this weak form of context grouping can be removed or relaxed further is an exciting direction for future work.\n\n**Q2:**  An essential reference missing is “Contextualize Me—The Case for Context in RL” by Benjamins et al., which discusses contextual RL extensively and also provides an online zero-shot analysis.\n\n**A2:**  Thank you for pointing out this important reference. We agree that the work by Benjamins et al. offers a valuable perspective on contextual RL and online zero-shot generalization. We will incorporate a discussion of this paper in the Related Work section to better situate our contributions and emphasize the broader significance of context in RL.\n\n**Q3:**  The notations are a bit heavy.\n\n**A3:**  Thank you for the helpful suggestion. We acknowledge that the variety of subscripts and superscripts—especially for multi-environment and time-step indexing—can become difficult to track. In the revision, we will include a notational summary table or add brief reminders in the text to improve readability. We also appreciate you pointing out minor typos, such as the phrase on line 316 (“... which we call it reinforcement ...”), and will correct them accordingly.\n\n**Q4:**  Same starting state is assumed. Is it necessary?\n\n**A4:**  Thank you for pointing this out. The assumption of a single starting state is made primarily for notational and explanatory simplicity, and it is without loss of generality. This assumption is common in the offline RL theory literature (e.g., [1]) and simplifies the presentation without affecting the core results. Our analysis can be extended to allow different initial states, as long as they are drawn from the same distribution. The main logic and proof techniques remain valid under this extension.\n\nTo avoid unnecessary technical clutter, we chose to adopt the single-start assumption in line with prior theoretical works. We will clarify this in the paper by adding a footnote, as per your helpful suggestion.\n\n[1] Provably Efficient Reinforcement Learning with Rich Observations via Latent State Decoding, ICML 2019.\n\n**Q5:**  Do the theorems place restrictions on the quality of the context?\n\n**A5:**  Thank you for the thoughtful question. Our theorems do not place any restrictions on the quality or informativeness of the context, beyond requiring that each context label $i$ is drawn i.i.d. from the same underlying context distribution $C$. The context labels need only distinguish environments from one another—they do not need to encode any semantic or structured information.\n\nEven if some environments differ significantly in difficulty or dynamics, the final policy is optimized to perform as well as possible on average across environments sampled from $C$. We will expand the discussion in the paper to clarify this point, following your helpful suggestion.\n\nAgain, thank you very much for your positive feedback and valuable suggestions. Please feel free to reach out with any further questions or comments.", "claims_and_evidence": "The claims in the paper are well supported by thorough and sound theoretical analysis.", "methods_and_evaluation_criteria": "N/A", "theoretical_claims": "I have checked the correctness of the proofs in the appendices A.1 and B.1 while only skimming the proofs in B.2 and C.\n\nOne point I remain uncertain about is assumptions about the context at hand. To me the work it seems tho implicitly assume \"perfect context information\". Under this assumption, the theoretical analysis makes perfect sense to me, though it would probably be best to highlight this assumption.", "experimental_designs_or_analyses": "N/A", "relation_to_broader_scientific_literature": "I believe the work puts itself well into the broader scientific literature and not only focusing on ZSG RL literature", "essential_references_not_discussed": "To me, one essential literature that is missing is \"Contextualize Me -- The case for context in RL\" (https://openreview.net/forum?id=Y42xVBQusn). This work discusses the contextual RL case at length and proved that optimal policies require context. Figure 1 of the work under review also reminded me of a similar figure in the paper by Benjamins et al. Though, this work views the problem of ZSG in RL from the classical online setting. Besides the theoretical component, that work proposes multiple contextual environments and shows the optimality gap for various different agent types empirically.", "other_strengths_and_weaknesses": "The work is generally well written, though with the variety of sub/super-scripts it can get difficult to properly track every detail.\nTo the best of my knowledge the work is the first providing a theoretical analysis of ZSG in offline RL.", "other_comments_or_suggestions": "Line 316: \"... which we call **it** *reinforcement* ...\" --> \"... which we call *reinforcement* ...\"", "questions_for_authors": "* In line131 it is stated that the same starting state is assumed for each $\\mathcal{M}_c$. Why is that needed?\n* Do your proofs require assumptions about the \"quality\" of the context or are they independent of these?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5S0YSUxlH2", "reviewer": "Reviewer_rcHz", "summary": "This paper investigates the generalization performance of policies learned in offline reinforcement learning (ORL) on test environments, specifically focusing on zero-shot performance in the average sense across test environments. The paper demonstrates that standard ORL methods, without access to contextual information, cannot possess the zero-shot generalization property (ZSG). To address this issue, the paper proposes a method based on contextual Markov Decision Processes (MDPs). This method leverages contextual information to partition the data into multiple MDPs and learns a policy that performs well on average across these MDPs. Theoretically, the paper proves that this method can learn a policy with ZSG.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks very much for reviewing our work and for the valuable advice. Our responses are as follows.\n\n**Q1:**  Differences from (Jin et al. 2021).\n\n**A1:**  We respectfully argue that our work is **substantially different** from (Jin et al., 2021), both in **algorithmic design** and **theoretical analysis**. The core challenge we address is *how to generalize across different environments*, rather than simply applying pessimism within each contextual MDP.\n- **Algorithmic Design:** While both works leverage the pessimism principle, we use it only as a **building block**. Specifically, Algo.1 performs pessimistic evaluation per environment, but this is only a **subroutine** for Algo.2 and Algo.3, which focus on **generalization across environments**. To our knowledge, this is the **first method** to adapt pessimism from time-step-wise to **environment-wise updates** for generalization. Jin et al. (2021), by contrast, operates within a single MDP and does not address this challenge.\n- **Theoretical Analysis:** In contrast to Jin et al. (2021), which analyzes a single MDP, we study the ZSG problem over multiple environments. For example, the proof of Theorem 9 (PERM) requires carefully separating SL and RL errors across environments (Eq. 22–28). Our analyses are non-trivial and require substantial innovation to handle generalization across distinct datasets. Additionally, Section 4 shows that the naive application of standard pessimistic offline RL in Jin et al. (2021) can yield **drastically suboptimal behavior**. This highlights that **simple extensions of existing pessimistic RL algorithms fail** in the ZSG setting.\n\n\n**Q2:**  About the context information:\n\n**A2:**  We would like to clarify that our algorithms do not require access to the context itself, nor does it rely on any semantic information. Instead, it only needs a **context indicator**—a label or identifier used to group trajectories that originate from the same environment, which often will be saved in the original dataset. In such cases, identifying the origin of each trajectory (i.e., grouping by environment) is often feasible. For this reason, we respectfully argue that this assumption is practical in many real-world scenarios.\n\n\n\nThat said, **relaxing the requirement** is an exciting future direction. One possibility is to infer latent environment clusters from trajectory data via representation learning which introduces challenges like cluster identifiability in the offline setting. We will clarify these points and expand the discussion in the final version, following your thoughtful suggestions.\n\n\n\n**Q3:**  No experiments.\n\n**A3:**  Thank you for the thoughtful suggestion. We would like to respectfully emphasize that this is a theoretical paper submitted to the **theory category**, and our primary goal is to develop first-principles theoretical insights into zero-shot generalization in offline RL—a significantly under-explored area compared to standard offline RL in a single environment or meta-RL approaches that rely on additional adaptation data. We believe that our theoretical contributions represent a substantial advancement in the foundations of reinforcement learning and are valuable in their own right. We plan to run a thorough experimental evaluation of our proposed PERM and PPPO algorithms in future work.\n\n**Q4:**  The relation to in-context or few-shot RL.\n\n**A4:**  In-context or few-shot RL typically assumes access to a small number of test-domain episodes or relies on history-dependent policies to enable **online adaptation**. In contrast, our setting focuses on zero-shot generalization, where **no additional interaction or adaptation** is allowed at test time. This necessitates learning a policy that performs well out of the box in any newly sampled environment from the context distribution.\n\n**Q5:**  If the MDPs are quite different, an optimal policy might differ from one MDP to another. How does one define zero-shot generalization meaningfully?\n\n**A5:**  In the ZSG setting, our goal is finding a single policy $\\pi^*$ that maximizes expected return across the distribution of environments:\n$\\pi^* = \\arg \\max_{\\pi} E_c[V_{c,1}^{\\pi}(x_1)]$. While $\\pi^*$ may not be optimal in every individual environment, it achieves the best average performance, which is **theoretically the strongest achievable guarantee without test-time context**. This objective follows prior work (Ye et al., 2023), which proves that achieving near-optimality in each environment (i.e., $E_c[V_c^{\\pi^*(c)}-V_c^{\\pi}]\\leq \\epsilon$) is impractical in the ZSG setting. We also note that this setup is meaningful and well-aligned with how generalization is defined in supervised learning: data points (contexts) are drawn from a distribution, and the goal is to learn a single predictor (policy) that performs well on average over that distribution. Allowing additional interaction with the environment makes our problem shift to meta-RL, which is beyond our ZSG scope.", "claims_and_evidence": "1. In my opinion, the proposed method looks like a straightforward extension on the previous work [Is pessimism provably efficient for offline rl?], i.e., it extends the pessimism to the contextual MDPs. The authors should provide more discussion about the difference between these two works (e.g., more theoretical properties and the behaviors of the learned policies), and do not be confined to the comparison for the theoretical bounds.\n\n2.  The context information should be given in the dataset, i.e., the dataset is divided by the contexts. This is hard to satisfy in pracetice. How will the algorithm behave in offline Contextual MDP if no context information is given?", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "n/a", "experimental_designs_or_analyses": "no experimtal results are provided.", "relation_to_broader_scientific_literature": "1. The considered situation is like the settings of domain generation where a few domains are given but no information about the new domain is given, and the task is to learn from the training domain a policy that is generalizable to the new domain. This paper addressed a similar problem under the offline RL settings and proposed two algorithms for it and provided theroretical analsis on that.\n\n2. The relation with in-context learning RL or a few shot learning RL should also be discussed.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "no experimental results provided, so that the feasibility and the effectiveness of the proposed algorithms are unkown.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "1. in line 157, the optimal policy is defined to be the one with highest expection value over the underlying MDPs, what about these MDPs are quite different? If they are different enough, an optimal policy may be different from MDP to MDP, which makes the average policy less useful. In other words, how to define the scope of generalization for zero shot learning? \n2. the need of contextual information seems to be a major limitation, how to relax that requirement?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The reviews of this paper illustrates the growing desire for applicability of ideas to actual data. Two reviewers weren't interested in the paper just because of its theoretical nature. 
Yet (qPEC) found it interesting to discuss because of some valid reasons:
- This is the first paper to provide theoretical analysis of ZSG in the offline RL setting (at least to our knowledge) and establishes a foundation for future work in this area.
- The paper is well-situated within broader RL literature
- The theoretical claims are carefully proven, with reviewer confirming the correctness of major results.

The key enabler for zero-shot generalization in this paper is having task identity labels (i.e., knowing which trajectories come from the same environment or task during training) so we are not doomed to do averages at test time. Yet at the core of the algorithm lies an Oracle for the uncertainty which reduces the applicability. The paper is showing what’s possible when pessimism + structured training are combined, and providing a template that can be extended to more realistic RL settings (and the Oracle is tractable for the linear case). 

Despite I regret that the algorithms are not tested (and it would be possible at a reasonable cost on toy tasks such as MDPs or different levels of a simple game), the work introduces an elegant new idea and I think is worth some attention.
Decision: Accept (poster)