Provable Zero-Shot Generalization in Offline Reinforcement Learning

Thanks very much for reviewing our work and for the valuable advice. Our responses are as follows.

Q1: Differences from (Jin et al. 2021).

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.

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

Q2: About the context information:

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.

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

Q3: No experiments.

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.

Q4: The relation to in-context or few-shot RL.

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.

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?

A5: In the ZSG setting, our goal is finding a single policy $\pi^*$ that maximizes expected return across the distribution of environments: $\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.

Thanks very much for your positive feedback and your valuable suggestions. Our responses are as follows.

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.

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.

Importantly, 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.

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

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.

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.

Q3: The notations are a bit heavy.

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.

Q4: Same starting state is assumed. Is it necessary?

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.

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

[1] Provably Efficient Reinforcement Learning with Rich Observations via Latent State Decoding, ICML 2019.

Q5: Do the theorems place restrictions on the quality of the context?

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.

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

Again, thank you very much for your positive feedback and valuable suggestions. Please feel free to reach out with any further questions or comments.

Thanks very much for reviewing our work and for the valuable suggestions. Our responses are as follows.

Q1: Augment theory with experiments; necessity of pessimism

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.

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

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

Alekh 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

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

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.

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

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?

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. We will clarify this point further in our revision to avoid any confusion following your advice.

Again, 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.