Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability

1 Introduction . Generalization is a central challenge in machine learning . However , much of the research on reinforcement learning ( RL ) has been concerned with the problem of optimization : how to master a specific task through online or logged interaction . Generalization to new test-time contexts has received comparatively less attention , although several works have observed empirically [ 1 , 2 , 3 , 4 ] that generalization to new situations poses a significant challenge to RL policies learned from a fixed training set of situations . In standard supervised learning , it is known that in the absence of distribution shift and with appropriate inductive biases , optimizing for performance on the training set ( i.e. , empirical risk minimization ) translates into good generalization performance . It is tempting to suppose that the generalization challenges in RL can be solved in the same manner as empirical risk minimization in supervised learning : when provided a training set of contexts , learn the optimal policy within these contexts and then use that policy in new contexts at test-time . Perhaps surprisingly , we show that such “ empirical risk minimization ” approaches can be sub-optimal for generalizing to new contexts in RL , even when these new contexts are drawn from the same distribution as the training contexts . As an anecdotal example of why this sub-optimality arises , imagine a robotic zookeeper for feeding otters that must be trained on some set of zoos . When placed ∗Equal contribution . 1 UC Berkeley , 2 Princeton University , 3 Facebook AI Research . Correspond to : dibya @ berkeley.edu , jrahme @ math.princeton.edu . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . in a new zoo , the robot must find and enter the otter enclosure . It can use one of two strategies : either peek through all the habitat windows looking for otters , which succeeds with 95 % probability in all zoos , or to follow an image of a hand-drawn map of the zoo that unambiguously identifies the otter enclosure , which will succeed as long as the agent is able to successfully parse the image . In every training zoo , the otters can be found more reliably using the image of the map , and so an agent trained to seek the optimal policy in the training zoos would learn a classifier to predict the identity of the otter enclosure from the map , and enter the predicted enclosure . This classification strategy is optimal on the training environments because the agent can learn to perfectly classify the training zoo maps , but it is sub-optimal for generalization , because the learned classifier will never be able to perfectly classify every new zoo map at test-time . Note that this task is not partially observed , because the map provides full state information even for a memoryless policy . However , if the learned map classifier succeeds on anything less than 95 % of new zoos at test-time , the strategy of peeking through the windows , although always sub-optimal in the training environments , turns out to be a more reliable strategy for finding the otter habitat in a new zoo , and results in higher expected returns at test-time . Although with enough training zoos , the zookeeper can learn a policy by solving the map classification problem , to generalize optimally when given a limited number of zoos requires a more intricate policy that is not learned by standard RL methods . How can we more generally describe the set of behaviors needed for a policy to generalize from a finite number of training contexts in the RL setting ? We make the observation that , even in fully-observable domains , the agent ’ s epistemic uncertainty renders the environment implicitly partially observed at test-time . In the zookeeper example , although the hand-drawn map provides the exact location of the otter enclosure ( and so the enclosure ’ s location is technically fully observed ) , the agent can not identify the true parameters of the map classifier from the small set of maps seen at training time , and so the location of the otters is implicitly obfuscated from the agent . We formalize this observation , and show that generalizing optimally at test-time corresponds to solving a partially-observed Markov decision process that we call an epistemic POMDP , induced by the agent ’ s epistemic uncertainty about the test environment . That uncertainty about MDP parameters can be modelled as a POMDP is well-studied in Bayesian RL when training and testing on a single task in an online setting , primarily in the context of exploration [ 5 , 6 , 7 , 8 ] . However , as we will discuss , this POMDP interpretation has significant consequences for the generalization problem in RL , where an agent can not collect more data online , and must instead learn a policy from a fixed set of training contexts that generalizes to new contexts at test-time . We show that standard RL methods that do not explicitly account for this implicit partial observability can be arbitrarily sub-optimal for test-time generalization in theory and in practice . The epistemic POMDP underscores the difficulty of the generalization problem in RL , as compared to supervised learning , and provides an avenue for understanding how we should approach generalization under the sequential nature and non-uniform reward structure of the RL setting . Maximizing expected return in an approximation of the epistemic POMDP emerges as a principled approach to learning policies that generalize well , and we propose LEEP , an algorithm that uses an ensemble of policies to approximately learn the Bayes-optimal policy for maximizing test-time performance . The primary contribution of this paper is to use Bayesian RL techniques to reframe generalization in RL as the problem of solving a partially observed Markov decision process , which we call the epistemic POMDP . The epistemic POMDP highlights the difficulty of generalizing well in RL , as compared to supervised learning . We demonstrate the practical failure modes of standard RL methods , which do not reason about this partial observability , and show that maximizing test-time performance may require algorithms to explicitly consider the agent ’ s epistemic uncertainty during training . Our work highlights the importance of not only finding ways to help neural networks in RL generalize better , but also on learning policies that degrade gracefully when the underlying neural network eventually does fail to generalize . Empirically , we demonstrate that LEEP , which maximizes return in an approximation to the epistemic POMDP , achieves significant gains in test-time performance over standard RL methods on several ProcGen benchmark tasks . 2 Related Work . Many empirical studies have demonstrated the tendency of RL algorithms to overfit significantly to their training environments [ 1 , 2 , 3 , 4 ] , and the more general increased difficulty of learning policies that generalize in RL as compared to seemingly similar supervised learning problems [ 9 , 10 , 11 , 12 ] . These empirical observations have led to a newfound interest in algorithms for generalization in RL , and the development of benchmark RL environments that focus on generalization to new contexts from a limited set of training contexts sharing a similar structure ( state and action spaces ) but possibly different dynamics and rewards [ 13 , 14 , 15 , 16 , 17 ] . Generalization in RL . Approaches for improving generalization in RL have fallen into two main categories : improving the ability of function approximators to generalize better with inductive biases , and incentivizing behaviors that are easier to generalize to unseen contexts . To improve the representations learned in RL , prior work has considered imitating environment dynamics [ 18 , 19 ] , seeking bisimulation relations [ 20 , 21 ] , and more generally , addressing representational challenges in the RL optimization process [ 22 , 23 ] . In image-based domains , inductive biases imposed via neural network design have also been proposed to improve robustness to certain factors of variation in the state [ 24 , 25 , 26 ] . The challenges with generalization in RL that we will describe in this paper stem from the deficiencies of MDP objectives , and can not be fully solved by choice of representations or functional inductive biases . In the latter category , one approach is domain randomization , varying environment parameters such as coefficients of friction or textures , to obtain behaviors that are effective across many candidate parameter settings [ 27 , 28 , 29 , 30 , 31 ] . Domain randomization sits within a class of methods that seek robust policies by injecting noise into the agent-environment loop , whether in the state [ 32 ] , the action ( e.g. , via max-entropy RL ) [ 14 ] , or intermediary layers of a neural network policy ( e.g. , through information bottlenecks ) [ 22 , 33 ] . In doing so , these methods effectively introduce partial observability into the problem ; while not necessarily equivalent to that of the epistemic POMDP , it may indicate why these methods generalize well empirically . Bayesian RL : Our work recasts generalization in RL within the Bayesian RL framework , the problem of acting optimally under a belief distribution over MDPs ( see Ghavamzadeh et al . [ 8 ] for a survey ) . Bayesian uncertainty has been studied in many sub-fields of RL [ 34 , 35 , 36 , 37 ] , the most prominent being for exploration and learning efficiently in the online RL setting . Bayes-optimal behavior in RL is often reduced to acting optimally in a POMDP , or equivalently , a belief-state MDP [ 6 ] , of which our epistemic POMDP is a specific instantiation . Learning the Bayes-optimal policy exactly is intractable in all but the simplest problems [ 38 , 39 ] , and many works in Bayesian RL have studied relaxations that remain asymptotically optimal for learning , for example with value of perfect information [ 5 , 40 ] or Thompson sampling [ 7 , 41 , 42 ] . Our main contribution is to revisit these classic ideas in the context of generalization for RL . We find that the POMDP interpretation of Bayesian RL [ 5 , 6 , 43 ] provides new insights on inadequacies of current algorithms used in practice , and explains why generalization in RL can be more challenging than in supervised learning . Being Bayesian in the generalization setting also requires new tools and algorithms beyond those classically studied in Bayesian RL , since test-time generalization is measured using regret over a single evaluation episode , instead of throughout an online training process . As a result , algorithms and policies that minimize short-term regret ( i.e. , are more exploitative ) are preferred over traditional algorithms like Thompson sampling that explore thoroughly to ensure asymptotic optimality at the cost of short-term regret . 3 Problem Setup . We consider the problem of learning RL policies given a set of training contexts that generalize well to new unseen contexts . This problem can be formalized in a Markov decision process ( MDP ) where the agent does not have full access to the MDP at training time , but only particular initial states or conditions . Before we describe what this means , we must describe the MDPM , which is given by a tuple ( S , A , r , T , ρ , γ ) , with state space S , action space A , Markovian transition function T ( st+1|st , at ) , bounded reward function r ( st , at ) , and initial state distribution ρ ( s0 ) . A policy π induces a discounted state distribution dπ ( s ) = ( 1−γ ) Eπ [ ∑ t≥0 γ t1 ( st = s ) ] , and achieves return JM ( π ) = Eπ [ ∑ t≥0 γ tr ( st , at ) ] in the MDP . Classical results establish that a deterministic Markovian ( memoryless ) policy π∗ maximizes this objective amongst all history-dependent policies . We focus on generalization in contextual MDPs where the agent is only trained on a training set of contexts , and seeks to generalize well to new contexts . A contextual MDP is an MDP in which the state can be decomposed as st = ( c , s′t ) , a context vector c ∈ C that remains constant throughout an episode , and a sub-state s′ ∈ S ′ that may vary : S : = C × S ′ . Each context vector corresponds to a different situation that the agent might be in , each with slightly different dynamics and rewards , but some shared structure across which an agent can generalize . During training , the agent is allowed to interact only within a sampled subset of contexts Ctrain ⊂ C. The generalization performance of the agent is measured by the return of the agent ’ s policy in the full contextual MDP J ( π ) , corresponding to expected performance when placed in potentially new contexts . While our examples and experiments will be in contextual MDPs , our theoretical results also apply to other RL generalization settings where the full MDP can not be inferred unambiguously from the data available during training , for example in offline reinforcement learning [ 44 ] .

Summary: The paper investigates the topic of generalization in reinforcement learning. It starts from the observation that even in fully observed environments, a kind of partial observability arises from epistemic uncertainty about the true MDP the agent faces at test time. This observation gives rise to what the paper calls an epistemic POMDP, which has the property that expected return in the epistemic POMDP equals test time agent performance in expectation over the posterior distribution of MDPs (given the prior is accurate).   This insight is used to point out that, being a POMDP, the optimal policy of the epistemic POMDP is in general memory-based, and in the class of memoryless policies, stochastic. In the same vein, even a policy that is optimal on all training contexts can have very poor generalization performance, and the optimal policy for the epistemic POMDP might not be optimal in any single MDP.  The remainder of the paper focuses on memoryless policies in the empirical epistemic POMDP, where the latter arises from a finite number of samples from the MDP posterior. It is shown that one can lower-bound the performance of an ensemble of policies on the empirical epistemic POMDP by their average performance minus a disagreement penalty to the combined policy \pi and that optimizing this bound gives the optimal policy for the empirical epistemic POMDP.  Based on these insights, the paper proposes a practical algorithm (LEEP) that combines bootstrap sampling from the training contexts, choosing max for combining the individual policies, PPO as the base algorithm, and the derived regularization term into an algorithm where the n individual policies are optimized in a round-robin fashion.  The proposed algorithm is evaluated on the procgen benchmark, where it compares favorably to pure PPO and Distral. The algorithm is also shown to avoid overfitting when only few training contexts are available. Finally, an ablation study shows that both the agreement penalty term and the max link function contribute to the performance advantage of LEEP.