When Is Generalizable Reinforcement Learning Tractable?

1 Introduction . Reinforcement learning ( RL ) is the dominant paradigm for sequential decision making in machine learning , and has achieved success in a variety of domains such as competitive gaming [ 33 , 40 ] and robotic control [ 21 , 22 ] . Despite this success , many issues prevent RL from being regularly used in the real world . For example , one typically trains and tests RL agents in the same environment . In such cases , an agent can memorize behavior that achieves high reward , without acquiring the true behavior that the system designer desires . This has raised concerns about RL agents overfitting to a single environment , instead of learning meaningful skills [ 17 ] . Indeed , a long line of work has noted the brittleness of RL agents : slight changes in the environment , such as those incurred by modeling or simulator design errors , or slight perturbations of the agent ’ s trajectory , can lead to catastrophic declines in performance [ 36 , 49 , 23 ] . Furthermore , although RL agents can solve difficult tasks , they struggle to transfer the skills they learned in one task to perform well in a different but similar task [ 37 , 48 ] . Yet , in the real world , it is reasonable to expect that RL agents will see scenarios that are at least mildly different from the specific scenarios they trained for . Hence , a desirable property of RL agents is that of generalization , broadly defined as the ability to discern the correct notion of behavior and perform well in semantically similar environments . We focus on two popular generalization settings . The Average Performance setting assumes there is an 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . underlying distribution over the environments that an agent might encounter . The agent ’ s goal is to perform well on average across this distribution [ 35 , 34 , 11 ] . The Meta Reinforcement Learning setting is closely related [ 20 , 10 , 37 ] . Here an agent first learns from a suite of training environments sampled from a distribution . Then at test time the agent must leverage this experience to adapt to a new environment sampled from the same distribution , via only a few queries in the new environment . Of course , in full generality , both notions of generalization are impossible to achieve efficiently . This is especially true in the RL function approximation setting , where the cardinality of the state space is potentially infinite , and so we desire query complexity that scales ( polynomially ) with the dimensionality of the state space [ 13 , 38 , 14 , 31 , 46 ] . Hence , key to both lines of inquiry is the premise that the environments are structurally similar . For example , a robot may face the differing tasks of screwing a bottle cap and turning a doorknob , but both tasks involve turning the wrist [ 37 ] . The hope is that if the environments are sufficiently similar , then RL can exploit this structure to efficiently discover policies that generalize . Yet , it remains unclear what kind of structure is necessary , and what it means for different environments to be close or similar . Motivated by this , we ask the following question : What are the structural conditions on the environments that permit efficient generalization ? This question underlies the analysis of our paper . We focus on environments that share state-action spaces , since even this basic case is not well understood in the literature . Indeed , even in this simplified setting , efficient generalization can be highly non-trivial . We make the following contributions . Our Contributions . We introduce Weak Proximity , a natural structural condition that is motivated by classical RL results , and requires the environments to have highly similar transition and reward functions and share optimal trajectories . We prove a statistical lower bound demonstrating that tractable generalization is impossible , despite this shared structure . This lower bound holds even when each individual environment can be efficiently solved to obtain an optimal linear policy , and when the agent possesses a generative model . Consequentially , we show that a classical metric for measuring the relative closeness of MDPs is not the right metric for modern RL generalization settings . Our lower bound implies that learning a state representation for the purpose of efficiently generalizing to multiple environments , is worst case sample inefficient — even when such a representation exists , the environments are ostensibly similar , and any single environment can be efficiently solved . To provide a sufficient condition for efficient generalization , we introduce Strong Proximity . This structural condition strengthens Weak Proximity by additionally constraining the environments to share an optimal policy . We provide an algorithm which exploits Strong Proximity to provably and efficiently generalize , when the environments share deterministic transitions . 2 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 [ 26 , 27 , 6 , 25 , 1 ] . 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 [ 32 , 7 , 24 , 43 ] 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 [ 18 ] , 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 [ 35 , 11 , 12 , 37 , 48 ] . On the theoretical side , [ 5 , 42 ] 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 [ 16 , 44 ] analyze the MAML algorithm [ 20 ] 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 [ 19 , 8 , 30 , 50 ] . Some works extend this to multiple MDPs [ 9 , 3 , 41 ] , 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 . 3 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 actionsA , 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 ⊂ Rd , 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 could have an uncountable state space . If we need to refer to the transition or reward function of a specific MDP M , then we shall denote this via T M or RM . We will denote a distribution over MDPs as D. We also assume that S can be partitioned into H different levels . This means that for each s ∈ S there exists a unique h ∈ { 0 , 1 . . . H − 1 } such that it takes h timesteps to arrive at s from s0 . We say that such a state s lies on level h , and denote Sh to be the set of states on level h. This assumption is without loss of generality , since we can always make the final coordinate of each state-action pair encode the number of timesteps that elapsed to reach the state . A “ deterministic MDP ” is one with deterministic transitions . For any MDP , we assume a single initial state s0 , which strengthens our lower bounds . A policy maps each state to a corresponding distribution over actions , and shall typically be denoted by π . The total expected reward accumulated by policy π when initialized at state s in MDP M is given by E [ ∑H−1 h=level ( s ) RM ( sh , ah ) | π ] and will be denoted by V sM ( π ) . Here the expectation is over the trajectory { ( sh , ah ) } H−1h=level ( s ) given that the first state in the trajectory is s. So V sM ( π ) is the value of the policy π in MDP M with respect to ( w.r.t ) initial state s. Analogously , if a policy is parameterized by θ = { θh } H−1h=0 , then we denote it as π ( θ ) , and the notation V sM ( π ) is then replaced by V sM ( θ ) . We assume that the cumulative reward collected by any trajectory from any initial state s in any MDP M is always bounded by 1 . Hence the value of any policy in any MDP lies in the interval [ 0 , 1 ] . TV ( P , Q ) denotes the total variation ( TV ) distance between probability distributions P and Q . 3.1 Problem Settings . Average Performance Setting . There is a fixed distribution D over a family of MDPs . One can sample MDPs from D. The algorithm can query states in the sampled MDPs , to learn some common structure . The goal is to solve max π EM∼D [ V s0M ( π ) ] . ( 1 ) Meta Reinforcement Learning Setting . There is a fixed distribution D over a family of MDPs . At training time , one can sample MDPs from D. The algorithm can query states in the sampled MDPs , to learn some common structure between all the MDPs . Then at test time , an MDP Mtest is sampled from the same distribution D. The goal of the algorithm is to learn a subroutine , which with non-trivial probability over the selection of Mtest , can solve max π V s0Mtest ( π ) , ( 2 ) significantly more efficiently than trying to solve Mtest without having seen any training MDPs . In both settings , “ sampling an MDP ” means drawing an MDP i.i.d from D , so that the agent can then interact with it by performing trajectories in it . Note that in Eqs . ( 1 ) & ( 2 ) , in full generality the initial state s0 is random and depends on M , Mtest . We focus on the case when the MDPs supporting D share a state-action space , and hence share the same single initial s0 since we assume a single initial state for any MDP . While such assumptions are already strong , they only strengthen our lower bounds . Furthermore , it is necessary to understand this simpler setting , before looking at more complex scenarios . To the best of our knowledge , such a study has not appeared in prior work . To solve the problems described by Eqs . ( 1 ) & ( 2 ) , we need to define an appropriate query model for the algorithm . We consider two query models , the first of which is strictly stronger than the second . Strong Query Model ( SQM ) . Sampling an MDP from D incurs no cost . The agent has a generative model of any sampled MDP M . To interact with M , the agent inputs a state-action pair ( s , a ) of M into the model , and receives RM ( s , a ) and a state sampled from T M ( s , a ) . This incurs a query cost of one . The goal is to solve Eqs . ( 1 ) & ( 2 ) with total query cost that is at most polynomial in |A| , H , d. Weak Query Model ( WQM ) . Sampling an MDP from D incurs a query cost of qD ≥ 1 . Within a sampled MDP M , the agent operates in the standard episodic RL setup . Concretely , during each episode the agent interacts with the MDP by starting from s0 , taking an action and observing the next state and reward , and repeating . Each action taken during an episode incurs a query cost of one . The goal is to solve Eqs . ( 1 ) & ( 2 ) with total query cost that is at most polynomial in qD , |A| , H , d. Note that under both SQM and WQM , we desire query cost that is polynomial in the dimension d of the state-action space , as opposed to the cardinality of the state space . This is standard for our function approximation setting [ 13 , 38 , 14 , 31 , 46 ] , since the cardinality of the state space could be infinite . Also , we separate SQM and WQM because it is well known that different query models can lead to various subtleties in analysis and sample complexity guarantees [ 13 , 38 , 14 , 31 , 46 ] . The generative model that defines SQM assumes that we can simulate any state of our choice without performing a trajectory , which is unrealistic in practice , and is one of the strongest oracle models considered in prior literature [ 28 , 4 , 39 , 14 , 31 , 2 ] . We shall present our lower bounds under SQM , which makes these results stronger , but shall present our upper bound under the natural and standard WQM . Without any conditions on D , the Average Performance & Meta RL settings can be intractable , even under SQM . This will occur if the MDPs supporting D do not share structure . This will also occur if any individual MDP can not be solved efficiently . Nevertheless , in practice one often deals with MDPs which share meaningful structure [ 11 , 37 ] . For instance , the transition distributions of the MDPs may be close in a suitable metric . Similarly , the reward functions of the MDPs might be close in an appropriate norm , or each MDP may share a set of optimal trajectories . And in practice , individual MDPs can usually be optimized efficiently [ 35 , 48 ] . In such cases , it is reasonable to expect tractable generalization . We are interested in formalizing conditions that permit efficient generalization . We will particularly focus on conditions which capture shared structure of the MDPs and the tractability of individual MDPs . We now formally state the problem we consider throughout our paper . Which conditions on D allow us to solve the Average Performance & Meta RL settings efficiently ? As mentioned above , there are two types of requirements . The first requirement should ensure that the MDPs are meaningfully similar . We formalize such conditions in Section 3.2 . The second requirement should ensure that any individual MDP is efficiently solvable , else there is no hope to efficiently find policies that generalize for many MDPs . We formalize such properties in Section 3.3 .

Summary: This paper studies the sample complexity of reinforcement learning in families of MDPs that share some notion of structure. The paper considers different notions of structure and sample complexity, and finds that even a relatively strong notion of structure where all MDPs in the family share an optimal linear policy and similar reward/transition functions, finding this shared optimal policy can require time exponential in the horizon of the problem. The paper also presents a stronger notion of structure under which finding a policy that attains near-optimal returns in all environments is tractable.