When Is Generalizable Reinforcement Learning Tractable?
Dhruv Malik, Yuanzhi Li, Pradeep Ravikumar
Abstract
Agents trained by reinforcement learning (RL) often fail to generalize beyond the environment they were trained in, even when presented with new scenarios that seem similar to the training environment. We study the query complexity required to train RL agents that generalize to multiple environments. Intuitively, tractable generalization is only possible when the environments are similar or close in some sense. To capture this, we introduce Weak Proximity, a natural structural condition that requires the environments to have highly similar transition and reward functions and share a policy providing optimal value. Despite such shared structure, we prove that tractable generalization is impossible in the worst case. This holds even when each individual environment can be efficiently solved to obtain an optimal linear policy, and when the agent possesses a generative model. Our lower bound applies to the more complex task of representation learning for the purpose of efficient generalization to multiple environments. On the positive side, we introduce Strong Proximity, a strengthened condition which we prove is sufficient for efficient generalization. Related Work Simulation Lemma. Many prior works define notions of statistical distance between Markov decision processes (MDPs), and measure the relative value of policies when deployed in different MDPs that are close under such metrics. The Simulation Lemma, which uses total variation distance between transitions and the absolute difference of rewards as this metric, is a well known formalization of this and has been very useful in classical prior work [KK99, KS02, BT03, KKL03, AN05]. These works do not directly tackle generalization, but their analyses construct an approximate MDP that models the true MDP under the aforementioned metric. Solving this approximate MDP then corresponds to solving the true MDP. It is natural to ask whether this metric is useful for measuring the similarity of MDPs in modern RL generalization settings. We show this metric is not appropriate for the settings we study. Transfer & Multitask Learning. There are varying formalisms of both settings, so we do not directly study them. However, they are broadly relevant, and we expect our theory to be useful for future studies of these settings. The works [LG10, BL13, Jia18, WZXS19] all study metrics for measuring variation between MDPs that are different from the metrics we study. A metric similar to the one used in the Simulation Lemma has also been studied [FYY19], and we show that this is inappropriate for our settings. Average Performance & Meta RL Settings. We directly study these two settings, which have seen much empirical work [PGK + 18, CKH + 19, CHHS20, RZF + 19, YQH + 19]. On the theoretical side, [BMPS20, SJT + 20] study an Average Performance setting where the agent receives a noisy observation in lieu of the actual state. We focus on the simpler setting where the agent knows its state. Recent works [FMO20, WCYW20] analyze the MAML algorithm [FAL17] in the context of Meta RL. In the worst case, their complexity bounds scale exponentially with the horizon, and they do not discuss structure which permits tractable Meta RL. Representation Learning. A large body of work has focused on extracting a representation useful for a single MDP [FPP04, Cas20, LBC21, ZMC + 21]. Some works extend this to multiple MDPs [CP10, AMCB21, SPM20], but they are about learning shared representations for MDPs that appear similar (but not from a sample efficiency perspective), while we formalize what it means for MDPs to be similar (in a sample efficient sense). Indeed, these works study the general case when the environments have distinct state spaces, but our lower bounds show generalization is non-trivial even when each MDP shares the same state space. Problem Formulation Notation & Preliminaries. Before describing our settings of interest, we establish notation and briefly review preliminaries. We always use M to denote a Markov decision process (MDP). Recall that an undiscounted finite horizon MDP is specified by a set of states S, a set of actions A, a transition function T which maps from state-action pairs to distributions over states, a reward function R which maps state-action pairs to nonnegative real numbers, and a finite planning horizon H. We assume that the state-action pairs are featurized, so that S × A ⊂ R d , and that (s, a) 2 = 1 for all (s, a) ∈ S × A. Any MDP we consider is undiscounted and has a finite action space, but
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext f8fa06ae-6053-46bf-ae6b-b3e0b505eb0cCited by top-tier papers8
- On the Importance of Exploration for Generalization in Reinforcement LearningYiding Jiang, J. Zico Kolter, Roberta RaileanuNeurIPS 2023 · 48 citations
- Overcoming the Sim-to-Real Gap: Leveraging Simulation to Learn to Explore for Real-World RLAndrew Wagenmaker, Kevin Huang, Liyiming Ke, Kevin Jamieson et al.NeurIPS 2024 · 45 citations
- Spectrum Random Masking for Generalization in Image-based Reinforcement LearningYangru Huang, Peixi Peng, Yifan Zhao, Guangyao Chen et al.NeurIPS 2022 · 33 citations
- Achieving Fairness Generalizability for Learning-based Congestion Control with JuryHan Tian, Xudong Liao, Decang Sun, Chaoliang Zeng et al.EuroSys 2025 · 10 citations
- On the Power of Pre-training for Generalization in RL: Provable Benefits and HardnessHaotian Ye, Xiaoyu Chen, Liwei Wang, Simon Shaolei DuICML 2023 · 8 citations
Builds on12
- Leveraging Procedural Generation to Benchmark Reinforcement LearningKarl Cobbe, Christopher Hesse, Jacob Hilton, John SchulmanICML 2020 · 685 citations
- Is a Good Representation Sufficient for Sample Efficient Reinforcement Learning?Simon S. Du, Sham M. Kakade, Ruosong Wang, Lin F. YangICLR 2020 · 213 citations
- Learning with Good Feature Representations in Bandits and in RL with a Generative ModelTor Lattimore, Csaba Szepesvári, Gellért WeiszICML 2020 · 181 citations
- Scalable Methods for Computing State Similarity in Deterministic Markov Decision ProcessesPablo Samuel CastroAAAI 2020 · 171 citations
- Observational Overfitting in Reinforcement LearningXingyou Song, Yiding Jiang, Stephen Tu, Yilun Du et al.ICLR 2020 · 148 citations
Related papers
- Learning Robust State Abstractions for Hidden-Parameter Block MDPsAmy Zhang, Shagun Sodhani, Khimya Khetarpal, Joelle PineauICLR 2021 · 5 citations
- A Generalized Bisimulation Metric of State Similarity between Markov Decision Processes: From Theoretical Propositions to ApplicationsZhenyu Tao, Wei Xu, Xiaohu YouNeurIPS 2025 · 6 citations
- Provably Efficient Causal Model-Based Reinforcement Learning for Systematic GeneralizationMirco Mutti, Riccardo De Santi, Emanuele Rossi, Juan Felipe Calderón et al.AAAI 2023 · 17 citations
- Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement LearningRishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, Marc G. BellemareICLR 2021 · 27 citations
- Representation Learning for Low-rank General-sum Markov GamesChengzhuo Ni, Yuda Song, Xuezhou Zhang, Zihan Ding et al.ICLR 2023
