Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability
Authors: Dibya Ghosh
Jad Rahme
Aviral Kumar
Amy Zhang
Ryan P Adams
Sergey Levine
Venue: NeurIPS 2021 Poster (2021)
Keywords: Reinforcement Learning
generalization

Abstract: Generalization is a central challenge for the deployment of reinforcement learning (RL) systems in the real world. In this paper, we show that the sequential structure of the RL problem necessitates new approaches to generalization beyond the well-studied techniques used in supervised learning. While supervised learning methods can generalize effectively without explicitly accounting for epistemic uncertainty, we describe why appropriate uncertainty handling can actually be essential in RL. We show that generalization to unseen test conditions from a limited number of training conditions induces a kind of implicit partial observability, effectively turning even fully-observed MDPs into POMDPs. Informed by this observation, we recast the problem of generalization in RL as solving the induced partially observed Markov decision process, which we call the epistemic POMDP. We demonstrate the failure modes of algorithms that do not appropriately handle this partial observability, and suggest a simple ensemble-based technique for approximately solving the partially observed problem. Empirically, we demonstrate that our simple algorithm derived from the epistemic POMDP achieves significant gains in generalization over current methods on the Procgen benchmark suite.


Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability
Dibya Ghosh :1
Jad Rahme +.2
Aviral Kumar!
Amy Zhang 1.3
Ryan P. 2dams
Sergey Levine'
Abstract
Generalization is central challenge for the deployment of reinforcement learning (RL) systems in the real world. In this paper, we show that the sequential structure of the RL problem necessitates new approaches generalization beyond the well-studied techniques used in supervised learning. While supervised learning methods can generalize effectively without explicitly accounting for epistemic uncertainty, we show that, perhaps surprisingly, this is not the case in RL. We show that generalization to unseen test conditions from limited number of training conditions induces implicit partial observability. effectively turning even fully- observed MDPs into POMDPs. Informed by this observation, recast the problem of generalization RL as solving the induced partially observed Markov decision process, which call the epistemic POMDP. We demonstrate the failure modes of algorithms that do not appropriately handle this partial observability, and suggest a simple ensemble-based technique for approximately solving the partially observed problem. Empirically, we demonstrate that our simple algorithm derived from the epistemic POMDP achieves significant gains in generalization current methods on the Procgen benchmark suite.
Introduction
Generalization is 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 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 new situations poses significant challenge to RL policies learned from 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 training set (i.e., empirical risk minimization) translates into good generalization performance. It is tempting to suppose that the generalization challenges in RL can solved the same manner as empirical risk minimization in supervised learning: when provided 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 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 of zoos. When placed
'Equal contribution. 1 UC Berkeley, Princeton University. 3 Facebook Al Research. Correspond to: waaame@math.ppiinettonccm dibya@berkeley. edu, jrahme@math, princeton.edd edu.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Strat Classify aandddraw Strat 2: Peek through habitat windows sucoeeds 10056 chance agent Is-able pars handwrtthen image
train time, pesking sub-optimal since learned classifer can get 100% acouraacy every
test time, peeking optimal
succonds with 25. chanco in every 200
tost accuracy of classier 96%,
by
Visualization of the robotic zookeeper example. Standard RL algorithms learn the classifier strategy, since it is optimal in every training zoo, but this strategy is sub-optimal for generalization because peeking generalizes better than the classifier at test-time. This failure occurs due to the following disconnect: while the task is fully--bserved since the image uniquely specifies the location of the otter habitat, to an agent that has limited training data, the location is implicitly partially observed at test-time because of the agent's epistemic uncertainty about the parameters of the image classifier.
Figure 1:
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 long 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 not partially observed, because the map provides full state information even for 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 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, 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 finite number of training contexts in the RL setting? We make the observation that, even in fully-observable domains, the agent's 's epistemic uncertainty renders the environment implicitly partially observed 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 cannot 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 implicitly obfuscated from the agent. We formalize this observation, and show that generalizing optimally at test-time corresponds to solving a partially call an epistemic POMDP, induced by the agent's epistemic uncertainty about the test environment.
Markov decision process that we
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 cannot collect more data online, and must instead learn 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-umin 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 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 's epistemic uncertainty during training. Our
2

---

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 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 RI., 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 cannot 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 class of methods that seek robust policies by injecting noise into the agent-envinon egent-envvronment loop, whether in the state [32]. the action (e.g. via max-entropp 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.
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 often reduced to acting optimally in POMDP, or equivalently, belief-state MDP [6], of which our epistemic POMDP is specific instantiation. Learning the Bayes-optimal policy exactly is intractable in all but the simplest problems [38, 391, 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 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 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.
Bayesian RL:
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 MDP M, which is given by tuple (S, (S, 7), with state space S, action space A. Markovian transition function T(*+4)/ 8t, at), bounded reward function r(st, and initial state distribution p(00). A
policy π induces a discounted state distribution dl (s) = (1 (1-110 7), 1r o ((s(s = s)], and achieves
3

---

Ideal Behavior
Learned RL Behavior
State
Action
Coat
T-Shirt
Boot
Coat
Coat
Coat
Coat
----1
X-=11
==0
-
X -1=11
X x--11
X r--1
t=1
t=2
tw
t=1
L=2
t=3
1=d
Episode ends
Sequential Classification RL Problem. In this task, an agent must keep guessing the label for an image until gets it correct. To avoid low test return, policies should change actions the label guessed was incorrect, but standard RL methods fail to do so, instead guessing the same incorrect label repeatedly.
Figure 2:
return (M( = Ex X (st, at)] in the MDP. Classical results establish that a deterministic Markovian (memoryless) policy x² maximizes this objective amongst all history-dependent policies.
We focus on generalization in contextual MDPs where the agent is only trained on 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 8, (c, s(), context vector e C that remains constant throughout an episode, and a sub-state E S' that may vary: S == x 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 sampled subset of contexts Ciran C C. The generalization performance of the agent is measured by the return of the agent's 's policy in the full contextual MDP J (T), 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 cannot be inferred unambiguously from the data available during training, for example in offline reinforcement learning [44].
4 Warmup: A Sequential Classification RL Problem
We begin our study of generalization in RL with an example problem that is set up to be close to a supervised learning task where generalization is relatively well understood: image classi- fication on the FashionMNIST dataset [45]. In this environment (visualized in Figure 2), an image from the dataset is sampled (the context) at the beginning of an episode and held fixed; the agent must identify the label of the image complete the episode. If the agent guesses correctly, it receives a reward of O and the episode ends; incorrect, it receives a reward of -1 and the episode continues, so it must attempt another guess for the same image the next time step. This RL problem is near identical to supervised classification, the core distinction being that an agent may interact with the same image over several timesteps in an episode instead of only one attempt as in super- vised learning. Note that since episodes may last longer than a single timestep, this problem is not contextual bandit.
Test Performance
Return
-10
0.0
0.5
1.0
1.5
2,0
# Gradient Steps
16b
DON (Test)
Uniform after step Adaptive
DQN (Train)
Number of attempts
Probabilty
0.8
0.7
0.1
The optimal policy in both the one-step and sequential RL sion of the problem deterministically outputs the correct label for the image, because the image fully determines the label (in other words, it is a fully observed MDP). However, this optimal strategy generally cannot be learned from finite training set, since some generalization error is unavoidable. With a fixed training set, the strategy for generalizing in classification re- mains the same: deterministically choose the label the agent is most confident about. However, the RL setting introduces two new factors: the agent gets multiple tries at classifying the same image, and it knows if an attempted label is incorrect. To generalize best to new test images, an RL policy must leverage this additional structure, for example by trying many possible labels, or by changing actions if the previous guess was incorrect.
5
* attempts before correct guess
Figure 3: DQN on RL FashionM- NIST. DQN achieves lower test perfor- mance than simple variants that lever- age the structure of the RI. problem.

---

A standard RL algorithm, which estimates the optimal policy on the empirical MDP defined by the dataset of training images does not learn to leverage these factors, and instead lears behavior highly sub-optimal for generalization. We obtained a policy by running DQN [46] (experimental details in Appendix A. whose policy deterministically chooses the same label for the image at every timestep. Determinism is not specific to DQN, and is inevitable in any RL method that models the problem as an MDP because the optimal policy in the MDP is always deterministic and Markovian. The learned deterministic policy either guesses the correct label immediately, or guesses incorrectly and proceeds to make the same incorrect guess on every subsequent time-step. We compare performance in Figure 3 with version of the agent that starts to guess randomly if incorrect on the first timestep, and a different agent that acts by process of elimination: first choosing the action it is most confident about, if incorrect, then the second, and forth. Although all three versions have the same training performance, the learned RI. policy generalizes more poorly than these alternative variants that exploit the sequential nature of the problem. In Section 5.2, we will see that this mocess-of-elimination is, in some sense, the optimal way to generalize for this task. This experiment reveals a tension: learning policies for generalization that rely on MDP model fail, even though the underlying environment is an MDP. This failure holds in any MDP model with limited data, whether the empirical MDP or more sophisticated MDPs that use uncertainty estimates in their construction.
5 Modeling Generalization in RL as an Epistemic POMDP
To better understand test-time generalization in RL, we study the problem under a Bayesian perspec- tive. We show that training on limited training contexts leads to implicit partial observability at test-time that we describe using a formalism called the epistemic POMDP.
5.1 The Epistemic POMDP
In the Bayesian framework, when learning given limited amount of evidence 'D from an MDP M. we can use a prior P(M) to construct a posterior belief distribution ((M(D over the identity of the MDP. For learning in contextual MDP, D corresponds to the environment dynamics and reward in training contexts Crainn that the agent can interact with, and the posterior belief distribution P(M) D) models the agent's uncertainty about the behavior of the environment in contexts that it has not seen before (c.g. uncertainty about the label for a test-set image in the example from Section 4).
Since the agent only has partial access to the MDP M during training, the agent does not know which MDP from the posterior distribution is the true environment, and must act at test-time under this uncertainty. Following reduction common in Bayesian RI. [6, 8), we model this test-time uncertainty using a partially observed MDP that we will call the epistemic POMDP. The epistemic POMDP is structured MDP M ~ P(M[D) from the posterior, and then the agent interacts with M until the episode ends in this MDP. The agent does not observe which MDP was sampled, and since the MDP remains fixed for the duration of the episode, this induces implicit partial observability. Effectively, each episode in the epistemic POMDP corresponds acting in one of the possible environments that is consistent with the evidence that the agent is allowed access to at training time.
follows: each new episode in the POMDP begins by sampling a single
The epistemic POMDP is formally defined as the tuple M (SP0, p A, TPO roo ppp .PP..)) 7)) A state in this POMDP x 8t) contains the identity of the current MDP being acted in M, and the current state in this MDP s;; we write the state space as S = M x S, where M is the space of MDPs with support under the prior. The agent only observes of = St. the state in the MDP (Ope S), but not the identity of the MDP, M. The initial state distribution is defined by the posterior distribution: ((M, so)) = P(M[D)) (8o), and the transition and reward functions in the POMDP reflect the dynamics in the current MDP:
( sa) (M, 5), = 8 (sMM(Mssssss..
= - (I)
=
Example (Sequential Image Classification). We begin by explicitly describing the induced epistemic POMDP for the task from Section 4. The agent's 's uncertainty concerns how images are mapped to labels, and each MDP M in the posterior distribution corresponds to different potential labelling function YM : x + y that is consistent with the training dataset. Each episode the epistemic POMDP, a different MDP M and corresponding labeller YM is sampled from the posterior distri- bution, alongside an image x ~ p(r). The agent must guess the label assigned by this labelling function y == YM(x), but is only provided the image and not the identity of the labeller YM. We
5

---

emphasize that the context remains fully observed in the epistemic POMDP (the image is provided the agent); what is partially observed is how the environment dynamics will behave for the context (what label the image corresponds to).
What makes the epistemic POMDP useful tool for understanding generalization in RL is that performance in the epistemic POMDP M corresponds exactly to the expected return of the agent test-time when the prior is well-specified.
Proposition 5.1. If the true MDP M is sampled from P(M), and evidence D from M is provided to an algorithm during training, then the expected test-time return of a is equal to its performance in the epistemic POMDP M
( = =BM-P(nn|UMM DD]
(2)
In particular, the optimal policy in M Bayes-optimal for generalization to the unknown MDP M: it receives the highest expected test-lime relurn amongst all possible policies.
The epistemic POMDP is based on well-understood concepts in Bayesian reinforcement learning, and Bayesian modeling more generally. However, in contrast to prior works on Bayesian RL, we are specifically concerned with settings where there a training- testt split, and performance is measured by single test episode. While using Bayesian RL to accelerate exploration or minimize regret has been well-explored [8], we rather use the Bayesian lens specifically to understand generalization perspective that is distinct from prior work on Bayesian RI. Towards this goal, the equivalence between rest-time return and expected return in the epistemic POMDP allows us to use performance in the POMDP as proxy for understanding how well current RL methods can generalize.
5.2 Understanding Optimality in the Epistemic POMDP
We now study the structure of the epistemic POMDP, and use it to characterize properties of Bayes- optimal test-time behavior and the sub-optimality of alternative policy learning approaches. The majority of our results follow from well-known results about POMDPs, we present them here informally, with formal statements and proofs in Appendix B.
Example ctd. Acting optimally in the epistemic POMDP for the sequential image classification task requires maximizing return over the distribution of labels that is induced by the posterior distribution p(y|a, Jr, D) = EM-P(M D) (Y = y)]. A deterministic policy (as is learned by standard RL algorithms) is high-risk strategy the POMDP: it receives exceedingly low return the labeller outputs a different label than the one predicted. The Bayes-optimal generalization strategy corresponds to a process of elimination: first choose the most likely label a = arg max p(y]ar, D); if this is incorrect, eliminate it and choose the next-most likely, repeating until the correct label is finally chosen. Amongst memoryless policies, the optimal behavior is stochastic, sampling actions according to the distribution x (alz) x /pyylx, D) (derivation in Appendix A.2).
The characteristies of the optimal policy in the epistemic POMDP for the image classification RL problem match well-established results that optimal POMDP policies are generally memory- based [47]. and amongst memoryless policies, the optimal policy may be stochastic [48, 49]. Because of the equivalence between the epistemic POMDP and test-time behavior, these maxims are also true for Bayes-optimal behavior when maximizing test-time performance.
Remark 5.1. The Bayes-optimal policy for maximizing test-time performance is in general non- Markovian. When restricted to Markovian policies, the Bayes-optimal policy is in general stochastic.
The reason that Bayes-optimal generalization often requires memory is that the experience collected thus far in the episode contains information about the identity of the MDP being acted in (which is hidden from the agent observation), and to maximize expected return, the agent must adapt its subsequent behavior to incorporate this new information. The fact that acting optimally at test-time formally requires adaptivity (or stochasticity for memoryless policies) highlights the difficulty of generalizing well in RL, and provides new perspective for understanding the success various empirical studies have found in improving generalization performance using recurrent networks [50, 51] and stochastic regularization penalties [32, 14, 22, 33].
is useful to understand what degree the partial observability plays a role in determining Bayes- optimal behavior. When the partial observability is insignificant, the epistemic POMDP objective can coincide with a surrogate MDP approximation, and Bayes-optimal solutions can attained with standard fully-observed RL algorithms. For example, if there is a policy that is simultancously

---

optimal in every MDP from the posterior, then an agent need not worry about the (hidden) identity of the MDP, and just follow this policy. Perhaps surprisingly, this kind of condition difficult relax: we show in Proposition B.I that even if a policy is optimal in many (but not all) of the MDPs from the posterior, this seemingly "optimal" policy can generalize poorly at test-time.
Moreover, under partial observability, optimal policies for the MDPs in the posterior may differ substantially from Bayes-optimal behavior: in Proposition B.2, we show that the Bayes-optimal policy may take actions that are sub-optimal in every environment in the posterior. These results indicate the brittleness of learning policies based on optimizing return in an MDP model when the agent has not yet fully resolved the uncertainty about the true MDP parameters.
Remark 5.2 (Failure of MDP-Optimal Policies, Propositions B.1, B.2). The expected test-time return of policies that are learned by maximizing reward in any MDP from the posterior, as standard RL methods do, may be arbilrarily low compared to that of Bayes-optimal behavior.
As Bayes-optimal memoryless policies are stochastic, one may wonder if simple strategies for inducing stochasticity, such as adding e-greedy noise or entropy regularization, can alleviate the sub- optimality that arose with deterministic policies in the previous paragraph. In some cases, this may be true; one particularly interesting result is that in certain goal-reaching problems, entropy-regue regula-reg RL can be interpreted as optimizing an epistemic POMDP objective with a specific form of posterior distribution over reward functions (Proposition B.3) [52]. For the more general setting, we show in Proposition B.4 that entropy regularization and other general-purpon catastrophically fail in epistemic POMDPs.
techniques can similarly
Remark 5.3 (Failure of Generic Stochasticity, Proposition B.4). The expected test-time return of policies learned with stochastic regularization techniques like maximun-entroop RL that are agnostic of the posterior P(M[D) may be arbitrarily low compared to that of Bayes-optimal behavior.
This failure happens because the degree of stochasticity used by the Bayes-optimal policy reflects the agent's epistemic uncertainty about the environment; since standard regularizations are agnostic to this uncertainty, the learned behaviors often do not reflect the appropriate level of stochasticity needed. A maze-solving agent acting Bayes-optimally, for example, may choose act deterministically in mazes like those it has seen at training, and on others where it is less confident, rely on random exploration to exit the maze, inimitable behavior by regularization techniques agnostic this uncertainty.
Our analysis of the epistemic POMDP highlights the difficulty of generalizing well in RL, in the complexity of Bayes-optimal policies (Remark 5.1) and the deficiencies of our standard MDP-based RL algorithms (Remark 5.2, 5.3). While MDP-based algorithms can serve as a useful starting point for acquiring generalizable skills, learning policies that perform well in new test-time scenarios may require more complex algorithms that attend to the epistemic POMDP structure that is implicitly induced by the agent's epistemic uncertainty.
6 Learning Policies that Generalize Well Using the Epistemic POMDP
When the epistemic POMDP M can be exactly obtained, we can learn RL policies that generalize well to the true (unknown) MDP M by learning an optimal policy in the POMDP. In this oracle setting, any POMDP (e.g. recurrent vs Markovian policies) or agent representations (e.g. belief state vs PSRs) made based on the requirements of the specific domain. However, in practice, the epistemic POMDP can be challenging to approximate due to the difficulties of learning coherent MDP models and maintaining a posterior over such MDP models in high-dimensional domains.
method will suffice, and design choices like policy function classes
In light of these challenges, we now focus on practical methods for learning generalizable policies when the exact posterior distribution (and therefore true epistemic POMDP) cannot be recovered exactly. We derive an algorithm for learning the optimal policy in the epistemic POMDP induced by an approximate posterior distribution P(M[D) with finite support. We use this to motivate LEEP, a simple ensemble-based algorithm for learning policies in the contextual MDP setting.
6.1 Policy Optimization in an Empirical Epistemic POMDP
Towards a tractable algorithm, we assume that instead of the true posterior P(M[D), we only have access to an empirical posterior distribution (MM[D) defined by n MDP samples from the posterior distribution [MM] This empirical posterior distribution induces empirical epistemic POMDP
7

---

Mr; our ambition is to learn the optimal policy in this POMDP. Rather than directly learning this optimal policy as a generic POMDP solver might, we recognize that M corresponds to collection of n MDPs and decompose the optimization problem to mimic this structure. We will learn policies -,,-. , TT, each policy in one of the MDPs M, from the empirical posterior, and combine these policies together to recover a single policy a for the POMDP. Reducing the POMDP policy learning problem into set of MDP policy learning problems can allow to leverage the many recent advances in deep RL for scalably solving MDPs. The following theorem links the expected return of a policy π in the empirical epistemic POMDP MMM in terms of the performance of the policies x; their respective MDPs M,,
Proposition 6.1. Let n,, "i,"".TT Tn be memoryless and define Tmax max4, ITM, (8,a)). The expected return of in MM is bounded below as:
1 n n
m2r max
J (T) 2 (m.)
E Bandi
[ VDKL (T((-1s) (-1ss)
(3)
71
This proposition indicates that if the policies in the collection [Tibseju) all achieve high return in their respective MDPs (first term) and are imitable by single policy æ (second term), then æ is guaranteed to achieve high return in the epistemic POMDP. In contrast, if the policies cannot be closely imitated by single policy, this collection of policies may not be useful for learning in the epistemic POMDP using the lower bound. This means that it may not sufficient to naively optimize each policy n. on its MDP M, without any consideration to the other policies MDPs, since the learned policies are likely to be different and difficult to jointly imitate. To be useful for the lower bound, each policy n should balance between maximizing performance on its MDP and minimizing its deviation from the other policies in the set. The following proposition shows that if the policies are trained jointly to ensure this balance, in fact recovers the optimal policy in the empirical epistemic POMDP.
Proposition 6.2. Let f [Ti]) + be a function that maps policies to single policy satisfying f(T, x) for every policy n, and let be a hyperparameter satisfying a (1-7)=n- letting ,i.... ni be the optimal solution the following optimization problem:
/2r Then
1 (m.) - -B [VDKL (Te(-s)) II {(I))(1))) (4)
{ heg[a] arg max "1:" Th
the policy n² = heem)) is optimal for the empirical epistemic POMDP M
6.2 A Practical Algorithm for Contextual MDPs: LEEP
Proposition 6.2 provides a foundation for a practical algorithm for learning policies when provided the proposition in
training contexts Curain from unknown contextual MDP. In order practical algorithm, we must discuss two problems: how posterior samples M; ~ P(M[D) can be approximated, and how the function f that combines policies should be chosen.
Approximating the posterior distribution: Rather than directly maintaining a posterior over transition dynamics and reward models, which is especially difficult with image-based observations, we can approximate samples from the posterior via a bootstrap sampling technique [53]. To sample a candidate MDP M,, we sample with replacement from the training contexts Crain to get new set of contexts Cinin- and define M, to be the empirical MDP on this subset of training contexts. Rolling out trials from the posterior sample M, then corresponds to selecting a context at random from Ca and then rolling out that context. Crucially, note that M, still corresponds to a distribution over contexts, not a single context, since our goal is to sample from the posterior entire contextual MDPs.
The link function ff in Proposition 6.2 that combines the set of poli- cies together effectively serves as an inductive bias: since we are optimizing in an approxima- tion to the true epistemic POMDP and policy optimization is not exact in practice, different choices can yield combined policies with different characteristics. Since optimal behavior in the epistemic POMDP must consider all actions, even those that are potentially sub-optimal in all MDPs in the posterior (as discussed in Section 5.2), we use an "optimistic" link function
Choosing a link function:
Note that when the true environment is a contextual MDP, the sampled MDP M, does not correspond to a single context within contextual MDP- each MDP M, is an entire contextual MDP with many contexts.
8

---

Algorithm Linked Ensembles for the Epistemic POMDP (LEEP)
1: Receive training contexts Crain number of ensemble members
2: Bootstrap sample training contexts to create C Crainn .. ..aiinn where Crain C Curain
3: Initialize n policies: 11,... .........
4: for iteration k = 1,2,3... 3,. do
5: for policy i 1, ,....:? do
6: Collect environment samples in training contexts Crain using policy ws
7:
Take gradient steps и, these samples with augmented RI. loss:
./Vi((ER (m)) ) + alR [PKI( max (y, (a(a)))))
Return π = max; x: T(as)) = max, x.(als)
8:
that does not dismiss any action that is considered by at least one of the policies, specifically f((Ti)E()) = (max, T.) = E. [s)
max (e
Algorithm: We learn a set of policies (EE)) using policy gradient algorithm implement the update step. To update the parameters for Tı, we take gradient steps via the surrogate loss used for the policy gradient, augmented by disagreement penalty between the policy and the combined policy f(( ((())))))) with a penalty parameter a > O, as in Equation 5:
[(m)) L (T,) C.E (ajsss max x (als)))
(5)
Combining these elements together leads to our method, LEEP, which we summarize in Algorithm 1. In our implementation, use PPO for CRT((T,)) [54]. In summary, LEEP bootstrap samples the training contexts to create overlapping sets of training contexts Crain Crain- Every iteration, each policy x; generates rollouts in training contexts chosen uniformly from its corresponding Ci Crain and then updated according to Equation 5, which both maximizes the expected reward and minimizes the disagreement penalty between each क. and the combined policy π max, T.
While this algorithm is structurally similar to algorithms for multi-task learning that train a separate policy for each context or group of contexts with a disagreement penalty [55, 56], the motivation and the interpretation of these approaches are completely distinct. In multi-task learning, the goal is to solve given set of tasks, and these methods promote transfer via loss that encourages the solutions to the tasks to be in agreement. In our setting, while we also receive set of tasks (contexts), the goal is not to maximize performance on the training tasks, but rather to learn a policy that maximizes performance on unseen test tasks. The method also has subtle but important distinction: each of our policies Ty acts on a sample from the contextual MDP posterior (which captures epistemic uncertainty), not single training context [55] or element from disjoint partitioning [56] (which does not). This distinction is crucial, since our generalization performance requires our aim is not to make it easier to solve the training contexts, but the opposite: prevent the algorithm from overfitting to the individual training contexts. Correspondingly, our experiments confirm that such multi-task learning approaches do not provide the same generalization benefits as approach.
7 Experiments
The primary ambition of our empirical study is to test the hypothesis that policies that are learned through (approximations of) the epistemic POMDP do in fact attain better test-time performance than those learned by standard RL algorithms. We do so on the Procgen benchmark [16]. challenging suite of diverse tasks with image-based observations testing generalization to unseen contexts.
1. Does LEEP derived from the epistemic POMDP lead to improved test-time performance over standard RL methods?
2. Can LEEP prevent overfitting when provided limited number of training contexts?
3. How do different algorithmic components of LEEP affect test-time performance ?
The Procgen benchmark a set of procedurally generated games, each with different generalization challenges. In each game, during training. the algorithm can interact with 200 training levels, before is asked to generalize to the full distribution of levels. The agent receives a 64 64 x image observation, and must output one of 15 possible actions. We instantiate method using ensemble

---

Maze
Heist
Bigfish
Dodgeball
LEEP
PPO
Env Steps 167
Env Steps lc7
# Env Steps
# Env Steps
Figure 4: Test set return for LEEP and PPO throughout training in four Proegen environments (averaged across random seeds). LEEP achieves higher test returns than PPO on three tasks (Maze, Heist and Dodgeball) and matches test return on Bigfish while having less variance across seeds.
of n 4 policies, penalty parameter of 1, and PPO [54] to train the individual policies (full implementation details in Appendix C).
We evaluate our method on four games in which prior work has found large gap between training and test performance, and which we therefore conclude pose a significant generaliza- tion challenge [16, 23, 26]: Maze, Heist, BigFish, and Dodgeball. In Figure 4, we compare the test-time performance of the policies learned using our method to those learned by a PPO agent with entropy regularization. In three of these environments (Maze, Heist, and Dodgeball), our method outperforms PPO by a significant margin, and in all cases, we find that the generalization gap between training and test performance is lower for method than PPO (Appendix D.I).
To understand how LEEP behaves with fewer training con-
texts, we ran on the Maze task with only 50 levels (Figure 5 (top)); the test return of the PPO policy decreases through training, leading to final performance worse than the starting random policy, but our method avoids this degradation.
Maze with 50 Levels
Train Return
10
Return
6
We perform ablation study the Maze and Heist envi- ronments (Maze in Figure 4, Heist in Appendix D.I) to rule potential confounding causes for the improved general- ization that our method displays on the Proegen benchmark tasks. First, see if the performance benefit derives solely from the use of ensembles, we compare LEEP to Bayesian model averaging strategy that trains an ensemble of policies without regularization ("Ensemble (no reg)"), and uses a mixture of these policies. This strategy does improve per- formance over the PPO policy, but does not match LEEP, indicating the usefulness of the regularization. Second, we compared to a version of LEEP that combines the ensemble policies together using the average 2 Σ (a ("IEEEP (avg)"). This link function achieves worse test-time perfor- mance than the optimistic version, which indicates that the inductive bias conferred by the max; link function is useful component of the algorithm. We also compare to Dis- tral, a multi-task learning method with different motivations but similar structure to LEEP: this method helps accelerate learning on the provided training contexts (figures in Appendix D. but does not improve generalization performance LEEP does. We additionally ablated the two key hyperparameters in LEEP, the number ensemble members n and the penalty coefficient (Table in Appendix D.2).
2 Stepsle7
Env Steps1e7
LEEP
PPO
Ablations on Procgen Mazes
7.0
6.0
5.5
4.5
4.0
0.0
0.5
1.0
1.5
2.0
2.5 107
Environment
LEEP
Ensemble ino rea)
Distral
LEEP (avg)
Figure 5: (top) Performance of LEEP and PPO with only 50 training levels Maze. (bottom) Ablations of I.EEP in Maze.
8 Discussion
It has often been observed experimentally that generalization in RL poses significant challenge, but it has so far remained an open question as to whether the RL setting itself presents additional generalization challenges beyond those seen in supervised learning. In this paper, we answer this question in the affirmative, and show that, in contrast to supervised learning, generalization in RL results in a new type of problem that cannot be solved with standard MDP solution methods, due to partial observability induced by epistemic uncertainty. We call the resulting partially observed setting the epistemic POMDP, where uncertainty about the true underlying MDP results challenging partially observed problem. We present practical approximate method that optimizes bound for performance an approximation of the epistemic POMDP, and show empirically that this approach, which we call LEEP, attains significant improvements in generalization over other RL methods that
10

---

do not properly incorporate the agent's 's epistemic uncertainty into policy optimization. A limitation of this approach that optimizes crude upproximation the epistemic POMDP with small number of posterior samples, and may be challenging to scale to better approximations to the true objective. Developing algorithms that better model the epistemic POMDP and optimize policies within is an exciting avenue for future work, and we hope that this direction will lead to further improvements in generalization in RL.
Acknowledgements
This research was supported by an NSF graduate fellowship, the DARPA assured autonomy program, the NSF IIS IIS-2007278 grant, a Princeton SEAS Innovation Grant and compute support from Google and Microsoft. We thank Benjamin Eysenbach. Xinyang Geng, and Justin Fu as well as members of the Princeton Laboratory for Intelligent Probabilistic Systems for helpful discussions and feedback.
References
[1] Jesse Farebrother, Marlos C. Machado, and Michael H. Bowling, Generalization and regulariza- tion in dqn. ArXiv, htss18800.00133, 2018.
[2] C. Zhang, Oriol Vinyals, R. Munos, and Bengio. study on overfitting in deep reinforcement learning. ArXiv, hts/1804406893, 2018.
[3] Niels Justesen, R. Torrado, Philip Bontrager, Ahmed Khalifa, J. Togelius, and S. Risi. Illuminat- ing generalization in deep reinforcement learning through procedural level generation. arXiv: Learning, 2018.
[4] Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfitting in reinforcement learning. ArXiv, abs/1912.00275, 2020.
[5] R. Dearden, N. Friedman, and Stuart Russell. Bayesian q-learning. In AAAI/IAII, 1998.
[6] M. Duff and Barto. Optimal learning: computational procedures for bayes-adaptive markov decision processes. 2002.
[7] M. Strens. A bayesian framework for reinforcement learning. In ICML, 2000.
[8] M. Ghavamzadeh, Shie Mannor, Joelle Pineau, and Aviv Tamar. Bayesian reinforcement learning: A survey. Found. Trends Mach. Learn., 8:359-483, 2015.
[9] A. Zhang, Yuxin Wu, and Joelle Pineau. Natural environment benchmarks for reinforcement learning. ArXiv, abs/ 101181166 .6002, 2018
[10] A. Zhang, Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization in continuous reinforcement learning. ArXiv, Abss/806.07937, 2018.
[11] S. Whiteson, B. Tanner, Matthew E. Taylor, and P. Stone. Protecting against evaluation overfitting in empirical reinforcement learning. 2011 IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning (ADPRL), pages 120-127, 2011.
[12] Zhuang Liu, Xuanlin Li, Bingyi Kang, and Trevor Darrell. Regularization matters in policy optimization an empirical study on continuous control. arXiv: Learning, 2020.
[13] Alex Nichol, Vicki Pfau, Christopher Hesse, Oleg Klimov, and John Schulman. Gotta learn fast:
A new benchmark for generalization in rl. ArXiv, abs/1804 1804.0372, 2018.
[14] K. Cobbe, Oleg Klimov, Christopher Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In ICML, 2019.
[15] Heinrich Kuttler, Nantas Nardelli, Alexander H. Miller, Roberta Raileanu, Marco Selvatici, Edward Grefenstette, and Tim Rocktäschel. The nethack learning environment. ArXiv, abs/2006. 13760, 2020.
[16] K. Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. ArXiv, ab5/1912.01888, 2020.
[17] Austin Stone, Oscar Ramirez, K. Konolige, and Rico Jonschkowski. The distracting control suite- a challenging benchmark for reinforcement learning from pixels. ArXiv, abs/2101.0272 2021.
11

---

[18] Max Jaderberg, V. Mnih, Wojciech Czarnecki, Tom Schaul, Joel Z. Leibo, D. Silver, and K. Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. ArXiv,
abs/161 1. 1.05997, 2017.
[19] Adam Stooke, Kimin Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. ArXiv, abs/2009.0833 2020.
[20] A. Zhang, Rowan McAllister, R. Calandra, Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. ArXiv, 0bs/2006.107422 2020.
[21] Rishabh Agarwal, Marlos C. Machaado P. S. Castro, and Marc G. Bellemare, Contrastive behav- ioral similarity embeddings for generalization in reinforcement learning. ArXiv, abs/2101.0526 2021.
[22] Maximilian Igl, K. Ciosek, Yingzhen Li, Sebastian Tschiatschek, C. Zhang, Sam Devlin, and Katja Hofmann. Generalization in reinforcement learning with selective noise injection and information bottleneck. In NeurlPS, 2019.
[23] Minqi Jiang. Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. ArXiv, abs/2010.03934, 2020.
[24] Kimin Lee, Kibok Lee, Jinwoo Shin, and H. Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. In ICI.R, 2020.
[25] Ilya Kostrikov, Denis Yarats, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. ArXiv, abs/2004. 13649, 2020.
[26] Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and R. Fergus. Automatic data augmentation for generalization in deep reinforcement learning. ArXiv, abs/2006. 12862,
2020.
[27] Fereshteh Sadeghi and Sergey Levine. (cad)2rl: Real single-image flight without single real image. ArXiv, abs/161 1.04201, 2017.
[28] Joshua Tobin, Rachel Fong, Alex Ray, J. Schneider, W. Zaremba, and P. Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 23-30, 2017.
[29] A. Rajeswaran, Sarvjeet Ghotra, Sergey Levine, and Balaraman Ravindran. Epopt: Learning robust neural network policies using model ensembles. ArXiv, abs/1610.01283. 2017.
[30] Xue Bin Peng. Marcin Andrychowicz, Wojeiech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pages 1-8, May 2018.
[31] Katie Kang, Suneel Belkhale, Gregory Kahn, Pieter Abbeel, and Sergey Levine. Generalization through simulation: Integrating simulated and real data into deep reinforcement learning for vision-based autonomous flight. In 2019 international conference robotics and automation (ICRA), pages 6008-6014, IEEE, 2019.
[32] F. Stulp, E. Theodorou, J. Buchli, and S. Schaal. Learning to grasp under uncertainty. 2011 IEEE International Conference on Robotics and Automation, pages 5703-5708, 2011.
[33] X. Lu, Kimin Lee, Abbeel, and Stas Tiomkin. Dynamics generalization via information bottleneck in deep reinforcement learning. ArXiv, abs/2008 000614 4, 2020.
[34] Deepak Ramachandran and E. Amir. Bayesian inverse reinforcement learning. In IJCAI, 2007.
[35] A. Lazaric and M. Ghavamzadeh. Bayesian multi-task reinforcement learning. In ICML, 2010.
[36] Wonseok Jeon, Seokin Seo, and Kee-Eung -ung Kim. A bayesian approach to generative adversarial imitation learning. In NeurIPS, 2018.
[37] Luisa M. Zintgraf, Shiarlis. Maximilian Igl, Sebastian Schulze, Y. Gal, Katja Hofmann, and S. Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta-learning. ArXiv, abs/1910.08348, 2020.
[38] R. Weber. On the gittins index for multiarmed bandits. Annals of Applied Probability, 2:1024 1033, 1992.
12

---

[39] P. Poupart, N. Vlassis, J. Hoey, and K. Regan. An analytic solution to discrete bayesian reinforcement learning. Proceedings of the 23rd international conference on Machine learning. 2006.
[40] R. Dearden, N. Friedman, and D. Andre. Model based bayesian exploration. In UAI, 1999.
[41] Ian Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. In NIPS, 2013.
[42] Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Math. Oper. Res., 39:1221112443 2014.
[43] Stephane Ross, Brahim Chaib-draa, and Joelle Pineau. Bayes-adaptive pomdps. In NIPS, pages 1225-1233 2007.
[44] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.00 006443,, 2020.
[45] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms, 2017.
[46] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Pe- tersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529-533, 2015.
[47] George E Monahan. State of the art-a survey of partially observable markov decision processes: theory, models, and algorithms. Management science, 28(1):1-16, 1982.
[48] Satinder P. Singh, Tommi S. Jaakkola, and Michael I. Jordan. Learning without state-estimation in partially observable markovian decision processes. In Proceedings of the Eleventh Inter- national Conference on International Conference on Machine Learning, page 284-292, San Francisco, CA, USA, 1994. Morgan Kaufmann Publishers Inc.
[49] Guido Montúfar, K. Zahedi, and N. Ay. Geometry and determinism of optimal stationary control in partially observable markov decision processes. ArXiv, abs/1503.0770 2015.
[50] Ilge Akkaya, Marcin Andrychowicz, Macick Chociej, Matcusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, al. Solving rubik's cube with a robot hand. Xiv preprint 1rXiv:1910.071133, 2019.
[51] Xue Bin Peng. Marcin Andrychowicz. Wojeiech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. CoRR, abs/1710.06553 2017.
[52] Benjamin Eysenbach and Sergey Levine. If maxent rl is the answer, what is the question? arXiv preprint arXiv:1910.01911 2019.
[53] Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
[54] John Schulman, Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv, t0s/ 000007003777744 2017.
[55] Yee Tch, Victor Bapst, Wojeiech M. Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
[56] Dibya Ghosh, Avi Singh, Aravind Rajeswaran, Vikash Kumar, and Sergey Levine. Divide-and- conquer reinforcement learning. In International Conference on Learning Representations, 2018.
[57] Scott Fujimoto, H. V. Hoof, and D. Meger. Addressing function approximation error in actor- critic methods. ArXiv, abs/1802.099777, 2018.
[58] Ilya Kostrikov. Pytorch implementations of reinforcement learning algorithms. https://
github. com/ http:/isstiittyooooooo
a2c- aww-ppa----pppppppppppaaaaaooaaa acktr-gail, 2018.
[59] Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning.
In In Proc. 19th International Conference Machine Learning. Citeseer, 2002.
13

---

[60] Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dy- namics model for generalization model based reinforcement learning, ArXiv, abs/2005.06800. 2020.
14


Reviews:
[{"review_id": "iXXGI92zg5", "reviewer": "Reviewer_pbCR", "summary": "The paper considers the problem of generalization in contextual MDPs, where the state of the environment consists of a context that remains fairly constant throughout an episode of the MDP. In such a setting the problem of generalization arises when an RL agent learns on a set of training MDPs (MDPs with contexts sampled from a set of training contexts) and then this trained RL agent must perform well on an unseen context test MDP. This paper underlines the challenges of learning and generalizing in such settings, namely, that methods that are not able to handle epistemic uncertainty fail at generalizing in such settings. The paper formulates this problem of generalization as a special POMDP and proposes a new sampling-based method for approximately solving this POMDP. Experimental results show that the proposed methods beats other baselines on a benchmark suite. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "Originality: To my knowledge the paper is original\n\nClarity: The paper is clear and well-written\n\nSignificance/Contributions: I think the paper makes significant theoretical, algorithmic and experimental contributions. It highlights the challenges of generalization for an RL agent and formulates the problem as a POMDP, then proposes a new method based on optimizing a proven lower bound on the original objective. The experiments are supportive of the method's ability to generalize on a benchmark suite. \n\nRelevance: Generalization in RL is an important problem to Neurips community and I think the paper is relevant to the Neurips community.\n\nComments: \n- The paper presents a very nice zookeeper example as motivation. Next, the paper also presents a very convincing RL as a classification task as an example. I was looking forward to see the results on the RL as a classification. I dont understand why the paper decided to experiment on a separate benchmark suite. It would have been interesting to see the performance of RL algorithm on the fashion MNIST classification task and to see if the proposed methods is able to come up with a policy that behaves as desired (that is after making an incorrect prediction first the agent changes its prediction) \nI think this is a major weakness of the paper since this exactly is the argument used by the paper to show that shortcoming of the current state of the art RL methods. Once this shortcoming is highlighted it is natural for the reader to expect that proposed method will address these shortcomings, but the paper fails to deliver on that. \n\n- Epistemic POMDP: What exactly are the key features of an epistemic POMDP that separates it from the a usual POMDP. Is it s sub-class of a general class of POMDPs or a super-class? A minor point here is that it seems that epistemic POMDP has an underlying policy where the agent predicts the hidden state of the world, specifically for the RL as a classification task. There already exists a class of POMDPs: POMDP with Information reward (and rhoPOMDP) that specifically address this issue. The paper might want to cite and contrast this sub-class of POMDPs. \n\n- I find it a bit surprising that in the experiments a very small value of n = 4 is used. In general, it is normal to use a large value of n for an ensemble method. It would have been nice to see how n matters for this method, even if not on a big deep RL task, but on a small task. \n\n- The paper also introduces a new parameter $alpha$ that balance the trade-off between generalization and overfitting on a context MDP. How to determine this hyperparameter. Experiments studying these hyperparameters are also missing. \n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The main limitation of the paper is that it highlights a weakness in the existing method, but then lacks experiments to support that motivation that the proposed method addresses these weaknesses. The paper currently is applied on a benchmark suite of games which I feel does not have a big negative societal impact. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "VXvpgxwqwm2", "reviewer": "Reviewer_Vmkh", "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). \n\nThis 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.\n\nThe 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.\n\nBased 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.\n\nThe 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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "First up, I would like to congratulate the authors on their clear writing. The paper is very easy to follow and especially the first half of the paper is a really enjoyable read. The topic of the paper, generalization, is important to the community and, as far as I know, the challenges faced in generalization settings have not been spelled out this clearly previously.\n\nThe paper is split into two parts - a motivational overview part up to and including page 6, while the remaining three pages propose a concrete algorithm based on theoretical insights and evaluate it in experiments.\n\nFor me the first half leaves little to be desired. Examples are clearly motivated and illustrative. Using a POMDP perspective to study generalization is an interesting perspective. My only nit-pick would be that the difference to Bayesian RL is over-emphasized. In particular in line L217 the paper seems to imply that Bayesian RL methods are not appropriate because they typically evaluate performance over multiple episodes. However, showing that a method performs well over multiple episodes does not imply that it doesn't on a single one.\n\nThe second part suffers as a consequence of the limited space and moves at a fast pace. One large concern is that it seems to me that Proposition 6.1 is over-interpreted. My understanding of (4) is that for any two policies $| J(\\pi) - J(\\pi_i) | \\leq \\alpha \\mathbb{E}[\\mathrm{KL}(\\pi, \\pi_i)]$. That is, any improvement that a policy $\\pi_i$ could have over $\\pi$ is smaller than the scaled, average KL. While this is a good motivation for trust-regions, in Proposition 6.1 this is exploited to conclude that by training on (5) we get a Bayes-optimal policy for the epistemic POMDP. However, from my understanding the only reason that we retrieve the optimal policy for the epistemic POMDP is that any improvement on $J_{M_i}(\\pi_i)$ by having $\\pi_i$ deviate from $f({\\pi_i})$ is necessarily smaller than the penalty term. As such, Proposition 6.1 does not make a statement about generalization, but a statement about the strength of the constraint and is largely equivalent to solving $\\mathrm{argmax}_{\\pi_i} sum J_{M_i}(\\pi_i)$ subject to the constraint $\\pi_i = f({\\pi_i})$. It would be great if the authors could comment on this as part of their rebuttal, since it is a key part of the motivation for LEEP.\n\nA critical component of LEEP seems to be the bootstrap sampling of the contexts. However, the discussion is limited to a single paragraph and there is no ablation study to determine its impact. It would be valueable to elaborate on the importance of this sampling and how the choice $n=4$ for the number of policies impacts experiments.\n\nIn general, the experiments focus mostly on a comparison with PPO. While this is in line with the main story about generalization, it would be really interesting how other methods cope with this setting to get an intuition of what the performance limits are here. E.g., including domain randomization (with access to more samples than the finite count) and stateful policies would be very valueable to determine the relative performance of LEEP and how much is (not) lost by focusing exclusively on static policies. After motivating stateful or stochastic policies, it would be great to see the tradeoff analyzed in more detail as part of the experiments. I also find it curious that a method like Distral performs worse, given its similarities to LEEP. Could you hypothesize about why this happens / what the crucial difference is? \n\n\n### Summary\n\nIn summary, I really like the presentation and perspective in the first half, but unfortunately this comes at a cost to the amount of information provided in the second part that proposes an algorithm. I would encourage the authors to trim down the first half (e.g., removing the otter example and focus on the example in Sec. 4) in order to give more room for explanations/interpretations in the second half of the paper. While I do have some concerns about the theoretical motivation and the depth of the experiments that will hopefully be addressed as part of the author feedback, I still think this could make a good paper for NeurIPS.\n\n\n### Minor & notational comments\n\n- There is some notational confusion between \\mathcal{M} and M. In the remainder of the notation sets/spaces are caligraphic, while instances are not, while for the MDP it is flipped. Also in 5.1 \\mathcal{M} is defined as a single MDP, while in 6.2 \\mathcal{M}_i is a distribution over MDPs. Being clear about what is meant where and being notationally consistent would make the paper more accessible.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are adequately presented, although it might make sense to be explicit in Proposition 6.1 that \"optimal\" means \"optimal memoryless\". \n\nAs a very general problem class, RL can be used for both benign and nefarious purposes. Better generalization would make RL better suited for both, but there don't seem to be any points that particularly emphasize the latter.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "4UwE-xvPi2z", "reviewer": "Reviewer_n4t5", "summary": "The paper addresses the problem of a reinforcement learning generalizing across different dynamics (with the same state/action space). Specifically, the paper frames this problem as a POMDP, where the properties of the environment are encoded as hidden information. An algorithm LEEP is introduced, which learns different policies for different sampled environments, regularized to make the policies similar, and then combines the policy into a single, hopefully robust policy.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "---Update After Author Response---\n\nMuch of my original review can be safely ignored, as I had a significant misunderstanding of the premise of the work. It is still included below for the sake of transparency. I have since carefully read the other reviews, the authors' responses to all reviews, and I have carefully re-read the paper.\n\nI now understand that the novel contribution of this paper is the development of a special case of Bayesian RL to the problem of generalizing to unseen states in an MDP. The main insight is that, when the agent has insufficient evidence to extrapolate its policy to unseen states, the agent effectively faces a partially observable problem where the partial observability results from the distribution of possible dynamics at the unknown state, based on the observed evidence. This is indeed, as far as I know, a clever and novel formulation of this issue. For instance, it happens to retroactively justify the common practice of random tie-breaking for action selection. In many typical implementations, if a state is completely unknown, with Q-values all 0, then the agent will act uniformly randomly, which this framework may very well recommend from a more principled position.\n\nBased on my initial confusion, I think I am qualified to suggest that the authors take great care in distinguishing their setting from the better studied setting of generalizing with unobserved context. Even in my re-read, I got tripped up in places and had to take some time to figure out how to reconcile what I was reading with what I understood to be the premise. I think one place where this is particularly important are in Section 5.1, where I think some bigger ballyhoo could be made of the fact that there is only one underlying MDP, that context is observable, and that this distribution is purely a result of the agent's beliefs about the dynamics of that one MDP. Another place where I got tripped up was in Section 4. For a bit I was not sure that the problem is actually fully observable -- I thought that perhaps the fact that the agent has already taken a particular action and received some feedback should be contained in the state description. It took me some time to convince myself that, in the traditional MDP formulation, this problem is simply multiple disconnected single state loops. Perhaps just a diagram of this MDP would make this more immediately clear to readers like myself?\n\nBeyond that, I think the current discussion between the authors and the other reviews covers other suggestions I would make.\n\n---Original Review Below---\n\n---Originality---\n\nThis is the area that I am most concerned about with this paper. There are two main contributions here: describing the dynamics generalization problem as a POMDP and the LEEP algorithm. Neither one of them seems particularly novel. \n\nWhile I can't off the top of my head point to a specific paper that explicitly formulates the generalization problem as a POMDP, it is my impression that it is pretty common to represent the \"context\" as hidden information. As one example, in \"Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning\" (Lee et al. 2020) they learn a dynamics model that attempts to extract hidden/latent contextual information from recent history and then to use that contextual information to inform decision-making and predictions. They cite prior work that operates on similar ideas and there may be follow-up work as well. They narrow their focus to this particular setting of hidden contextual information, and give that special case its own name (CaDM) rather than casting it in the much broader POMDP framework, but it seems to me that the insight gain from the presentation in this paper is minimal.\n\nThe paper seems to suggest that casting this problem as a POMDP offers novel insight into how difficult the dynamics generalization problem is. It even goes so far as to say that there was an \"open question\" about whether it is difficult to generalize to new dynamics in RL and that this paper provides the answer. I don't find that to be the case. Again, existing work addressing this problem already treats context as hidden state information and already learns policies that are history dependent or stochastic, so noting that this may be necessary seems like old news.\n\nFurthermore, in this paper the formulation as a POMDP doesn't seem to actually go very far. If I am understanding correctly, LEEP is still learning a memoryless policy and does not do any of the Bayesian reasoning about the hidden information that a solution to the POMDP would do. Existing work (e.g. that Lee et al. paper above) seems to already attack this in a more sophisticated way, where the policy can react and change when recent observations offer evidence about what context it might be operating in.\n\nLEEP itself is, as the paper acknowledges, very similar to existing multi-task RL and meta-RL approaches. In a vacuum, I think LEEP does not really add much to what is already out there. The paper suggests that LEEP is notable because it is motivated by this POMDP formulation, but, again, that connection seems tenuous at best.\n\nSo I am not sure what this paper substantively adds to the body of work already out there on dynamics generalization in RL. If I am underestimating the novel contribution, then I think the paper needs to be far more explicit about contrasting this formulation/approach with similar, existing ideas.\n\n---Quality---\n\nI don't detect any technical flaws in the theoretical work and the experiments seem well-designed.\n\nAs described above, I do think the connection between the POMDP formulation and the algorithm itself is overstated. I don't really see where in this algorithm POMDP-ish inference is supposed to be occurring. For instance, the way the final policy is construction from individual policies is not really ever justified, especially not as a policy that would be suitable for this POMDP.\n\nThe experiments do a good job of empirically answering the questions posed. However, a big remaining empirical and conceptual question mark is what are LEEP's relative strengths and weaknesses, compared to existing approaches that aim to improve dynamics generalization in RL?\n\n---Clarity---\n\nI found the paper to be well-written and clearly presented.\n\n---Significance---\n\nGiven that I don't see a lot of novel insight in this paper, and given that the algorithm is not adequately contextualized in the related literature on this problem, I would say that the likely impact of this paper would be low.\n\n---Summary of the Review---\n\nThe paper is well-written and technically sound, but it seems to me that its main contributions and insights are not sufficiently novel with respect to work already published in this area.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper does not discuss negative societal impact, but I do not expect this work to have significant societal impact in the foreseeable future.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "3M2iNSzdpRA", "reviewer": "Reviewer_C295", "summary": "* This paper uses the insights from Bayesian RL to solve the problem of generalization in RL. In particular, the authors formulated the generalization problem as an epistemic POMDP and proposed an ensemble-based algorithm, called LEEP, to approximately solve the problem.\n\n* The authors empirically demonstrated the performance of LEEP in generalizing to unseen contexts in the Procgen benchmark. They have shown that:\n  * LEEP performed better than PPO in Maze, Heist, BigFish, and Dodgeball.\n  * An ablation study to show that the better generalization performance of LEEP does not only come from the use of ensembles.\n  * An ablation study to show that the better generalization performance of LEEP does come from the inductive bias conferred by the `max_i π_i` link function proposed by the authors.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "", "main_review": "* This paper proposes a novel problem: the problem of generalization in RL.\n  * This problem is closely related to the problem of *Bayesian RL*.\n  * However, the difference is that in the authors' problem, the RL agent is **not** allowed to update its belief / posterior after collecting new data. By contrast, a Bayesian RL agent updates its belief / posterior at every iteration.\n  * Therefore, Bayesian RL focuses on improving the performance in unseen situations by optimally trading off between exploration vs. exploitation. In contrast, the authors' problem focuses on improving the performance in unseen situations by **exploitation only**.\n  * Based on this problem definition, the authors proposed an algorithm to train a policy in a certain way such that this policy can do well in unseen situations even **without any exploration**.\n\n\n* I think the problem formulation and the algorithm are sound.\n  * In line 105, the authors briefly discuss the similarity between domain randomization and the proposed approach. It could make the paper stronger to compare the proposed algorithm with domain randomization in the Procgen benchmark.\n  * In line 78, the author says that `Our work highlights the importance of not only (1) finding ways to help neural networks in RL generalize better, but also (2) on learning policies that degrade gracefully when the underlying neural network eventually does fail to generalize.`\n    * I am not sure about the difference between (1) and (2). If they are indeed different, has (2) been demonstrated in the results? (it is clear that (1) has been shown in the results).\n\n\n\n* I have not seen many other works solving this problem.\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "* I was very confused at the beginning about the difference between this paper and Bayesian RL. Assuming that my understanding of this paper is correct, I think the writing could be improved. Here are my comments:\n  * It might be helpful to have a clear paragraph about the `problem definition` of this paper. I have found it confusing between Bayesian RL and the problem that this paper is focusing on.\n  * Line 29-45: the `zoo` example conveys the drawback of empirical risk minimization. I understand that the `map` represents empirical risk minimization. However, I am not sure what `peeking through the window` represents? The `anti-empirical risk minimization approach`?\n  * Line 29-45: the `zoo` example is good. However, it is an analogy for supervised learning rather than an analogy for RL. Considering that this paper is doing RL, it might be better to use an analogy for RL.\n  * Sec. 4 and Fig. 1: the authors try to explain why standard RL methods could fail to generalize. However, the example problem of shoe classification is a contextual bandit problem rather than an MDP problem. Since the authors mentioned `MDP` multiple times in Sec. 4, it might be better to give an example of an MDP rather than a bandit.\n  * Sec. 4 and Fig. 1: When I first read Sec. 4 and Fig. 1, I was very confused.\n    * Given that methods like UCB can solve bandit problems, when I first read it, I didn't understand why the author said `more sophisticated MDPs that use uncertainty estimates in their construction.` Later on, I understood that the authors are solving a different problem from the problem that UCB is trying to solve.\n    * Therefore, I think this section could be clarified a bit to avoid this confusion.\n  * Line 228-237: again, I think this is a bandit problem, not an MDP problem.\n  * Line 228-237: in this paragraph, the authors use terms like `Bayes-optimal,` which is very confusing to me at first.\n    * `Bayes-optimal` refers to optimally solving the trade-off between exploration and exploitation. However, in the problem that this paper is focusing on, there is no exploration.\n      * It is very confusing to see `bayes-optimal` and `a = argmax p(y | x, D);` together. Because in Bayesian RL and bandits, `a = argmax p(...)` usually refers to a greedy heuristic to **approximate** the bayes-optimal policy, e.g., UCB. So `a = argmax p(...)` is **not** bayes-optimal.\n    * I think maybe the author could define a new concept of optimality for the generalization problem formulated in this paper, rather than using terms in Bayesian RL.\n  * Line 228-237: again, the term `POMDP` makes me confusing at first, too.\n    * In a POMDP, the belief (or posterior) is updated at every iteration when a new observation comes. However, in the problem formulated in this paper, no exploration (active information gathering) is allowed, and the belief is not updated at all during testing. Thus, I am not sure whether `POMDP` is the correct term. Maybe the authors could use `POMDP` and clarify this point.\n  * Line 236: the authors use the word `memoryless` without defining it. I figured it out by reading the appendix. It would be better to define it in advance.\n  * Eq. 6: is the policy gradient in Eq. 6 solving the optimal problem? So after convergence, will we get the optimal solution to Eq. 5? It might be better to clarify.\n\n\n* Minor\n  * Line 78: `but also on learning` - `on` is unnecessary.\n  * Line 132: `dπ(s) = (1−γ)...` - I don't understand why `(1−γ)` is there (I could be wrong).\n  * Line 212: extra space in the beginning.\n  * Line 216: extra space in the beginning.\n  * Line 233: It might be better to use a larger symbol for the indicator function.\n\n* I didn't check the proofs in Appendix.\n\n\n# Post-rebuttal\nThank you for your clarification! Now I think the paper is clearer in my mind and I appreciate it more! I have raised my score by 1 point.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper provides a conceptually interesting perspective on dynamics generalization in RL by drawing a connection to a special subtype of the Bayesian RL formalism. The work's experiments got sufficiently beefed as a result of the discussion with the reviewers, and now the paper is strong both from the conceptual and empirical standpoints. It is important that the authors pay special attention to incorporating the clarifications they gave in response to reviewers' questions into the final paper version. The paper's impact will heavily depend on its clarity and its success in ensuring that readers will see the distinction of the paper's model from general Bayesian RL, something that initially confused several reviewers.
Decision: Accept (Poster)