Explore to Generalize in Zero-Shot RL
Authors: Ev Zisselman
Itai Lavie
Daniel Soudry
Aviv Tamar
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement Learning
Generalization
State Space Maximum Entropy Exploration
TL;DR: We study zero-shot generalization in reinforcement learning, i.e., we aim to find a policy on a set of training tasks such that it will perform well on a similar, but unseen test task.

Abstract: We study zero-shot generalization in reinforcement learning - optimizing a policy on a set of training tasks to perform well on a similar but unseen test task. 
To mitigate overfitting, previous work explored different notions of invariance to the task. However, on problems such as the ProcGen Maze, an adequate solution that is invariant to the task visualization does not exist, and therefore invariance-based approaches fail. 
Our insight is that learning a policy that effectively $\textit{explores}$ the domain is harder to memorize than a policy that maximizes reward for a specific task, and therefore we expect such learned behavior to generalize well; we indeed demonstrate this empirically on several domains that are difficult for invariance-based approaches. Our $\textit{Explore to Generalize}$ algorithm (ExpGen) builds on this insight: we train an additional ensemble of agents that optimize reward. At test time, either the ensemble agrees on an action, and we generalize well, or we take exploratory actions, which generalize well and drive us to a novel part of the state space, where the ensemble may potentially agree again. We show that our approach is the state-of-the-art on tasks of the ProcGen challenge that have thus far eluded effective generalization, yielding a success rate of 83% on the Maze task and 74% on Heist with $200$ training levels. ExpGen can also be combined with an invariance based approach to gain the best of both worlds, setting new state-of-the-art results on ProcGen.
Code available at [https://github.com/EvZissel/expgen](https://github.com/EvZissel/expgen).

Explore to Generalize in Zero-Shot RL
Ev Zisselman; Itai Lavie, Daniel Soudry, Aviv Tamar Technion Israel Institute of Technology
Abstract
We study zero-shot generalization reinforcement Barning-optimiziing policy on a set of training tasks to perform well on a similar but unseen test task. To mitigate overfitting, previous work explored different notions of invariance to the task. However, on problems such as the ProeGen Maze, an adequate solution that is invariant to the task visualization does not exist, and therefore invariance-base approaches fail. Our insight is that learning policy that effectively explores the domain is harder to memorize than a policy that maximizes reward for a specific task, and therefore expect such learned behavior generalize well; indeed demonstrate this empirically on several domains that are difficult for invariance- based upproaches. Our Explore to Generalize algorithm (ExpGen) builds this insight: we train an additional ensemble of agents that optimize reward. At test time, either the ensemble agrees on action, and generalize well, exploratory actions, which generalize well and drive us to a novel part of the state space, where the ensemble may potentially agree again. We show that approach is the state-of-the on tasks of the ProcGen challenge that have thus far eluded effective generalization, yielding success rate of 83% the Maze task and 74% on Heist with 200 training levels. ExpGen can also be combined with an invariance based approach to gain the best of both worlds, setting new state of of-the-ard art results on ProcGen. Code available at https: //gi///ii /ightut...... cww Coom/Evisselll exxpgen
take
1 Introduction
Recent developments in reinforcement learning (RL) led to algorithms that surpass human experts in a broad range of tasks [Mnih et al. 2015 Vinyals al. 2019 Schritwieser et al. 2020 Wurman et al.) 2022]. In most cases, the RL agent is tested on the same task was trained on, and is not guaranteed perform well unseen tasks. In zero-shot generalization for RI (ZSG-RL)) however, the goal is to train an agent on training domains to act optimally in a new, previously unseen test environment Kirk al.] 2021). A standard evaluation suite for ZSG-RL is the ProcGen benchmark (Cobbe et al.; 2020). containing 16 games, each with levels that are procedurally generated to vary in visual properties (e.g., color of agents in BigFish, Fig. Ta background image in Jumper, Fig. Ic] and dynamics (e.g..
wall positions in Maze, Fig. ld and key positions in Heist, Fig. Tep.
Previous studies focused identifying various invariance properties in the tasks, and designing corresponding invariant policies, through an assortment of regularization and augmentation tech-
niques [gl et al.; 2019 Cobbe et al. 2019 Wang al. [2200] Lee et al. 2019a Raileanu et al. 2021
Raileanu and Fergus] 2021 Cobbe et al. [202 [Sonar et 2021 Bertran et al. 2020 [i et al. 2021
For example, policy that invariant to the color of agents is likely generalize well in Big Fish.
More intricate invariances include the order of observations
a trajectory [Raileanu and Fergus
2021]. and the length of a trajectory, reflected in the value function [Raileanu and Fergus 2021
"Conrespondence E-mail: ev ev_zis@campus. technion.cc.... ac. ac.il
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

EC
(a) BigFish
(b) Plunder
(c) Jumper
(d) Maze
(c) Heist
(f) Hidden Maze
Figure 1: (c))( and (e) displays screenshot of ProcGen games. (f) Imaginary maze with goal
and walls removed (see text for explanation).
Can ZSG-RL. be reduced to only finding invariant policies? As a counter-argument consider the following thought experimen + Imagine Maze, but with the walls and goal hidden in the observation (Fig. lf.. Arguably, this is the most task-invariant observation possible, such that a solution can still be obtained in reasonable time. An agent with memory can be trained to optimally solve all training tasks: figuring out wall positions by trying to move ahead and observing the resulting motion, and identifying based on its movement history which training maze it is currently in. Obviously, such strategy will not generalize to test mazes. Indeed, as depicted in Figure 2 performance in tasks like Maze and Heist, where the strategy for solving any particular training task must be indicative of that task, has largely not improved by methods based on invariance (e.g. UCB-DrAC and IDAAC).
Interestingly, decent zero generalization can be obtained even without a policy that generalizes well. As described by Ghosh et al. [2021], an agent can overcome test-time errors in its policy by treating the perfect policy as unobserved variable. The resulting decision making problem, termed the epistemic POMDP, may require some exploration at test time to resolve uncertainty. Ghosh et al. [202T] further proposed the LEEP algorithm based on this principle, which trains an ensemble of agents and essentially chooses randomly between the members when the ensemble does not agree, and was the first method to present substantial generalization improvement on Maze.
In this work, we follow the epistemic POMDP idea, but ask: how to improve exploration at test time? Our approach is based on a novel discovery: when train an agent to explore the training domains using a maximum entropy objective [Hzaan al. 2019 Mutti et al. 202 we observe that the learned exploration behav- ior generalizes surprisingly well well-much better than the generalization attained when training the agent maximize reward. Intuitively, this can be explained by the fact that reward is a strong signal that leads to specific behavior that the agent can 'memorize" during training, while exploration is naturally more varied, mak- ing it harder to memorize and overfit.
Norma lized performance ProcGen games
0.8
0.6
0.4
BigFish
Plunder
Jumper
Maze
Heist
Figure 2: Normalized test Performance for Exp- Gen, LEEP, IDAAC, DAAC, and PPO, five ProcGen games. ExpGen shows state-of-the-as performance test levels of Maze, Heist and Jumper; games that are notoriously challenging for other leading approaches. The scores nor- malized as proposed by [Cobbe et al. 2020
Exploration by itself, however, is not useful for solving new tasks. Our algorithm, Explore to Generalize (ExpGen), additionally trains an semble of either the ensemble agrees generalize well, or we take exploratory actions using the exploration policy, which we demon- strate to generalize, and drive us to a novel part of the state space, where the ensemble may potentially agree again.
agents. At test time,
action, and
ExpGen simple to implement, and can used with any reward maximizing RL algorithm. Combined with vanilla PPO, ExpGen significantly improves the state-of-the of-the-art (SOTA) on several ProcGen games for which previous methods fail (see Fig. 2}. ExpGen also significantly improves
We validated this experiment empirically, using recurrent policy on Maze with 128 training tasks, observing 85% success rate training domains, and test success similar random policy (see Appendix(A)

---

upon LEEP, due to its effective test-time exploration strategy. For example, on Maze with 200 training levels, our method obtains 83% success on test tasks, whereas the previous oftt--o-thhearr achieved 66%. When combined with IDAAC [Raileanu and Fergus, 2021]. the leading invariance- based algorithm, ExpGen achieves state-of-the-ar performance on the full ProcGen suite (the full results are provided in Appendix D
2 Related Work
Generalization in RL The recent survey by Kirk et al.
[021 provides an extensive review of
generalization RL; here, we provide a brief overview. One approach to generalization is by
artificially increasing the number of training tasks, using either procedural generation [Cobbe et al.
2019 2020]. or augmentations [Kostrikov et al. 2020, Ye et al. 2202 Lee et al. 2019a Raileanu et al.
2021
task interpolation (Yao et al. 2021] various regularization technique, such as dropout [Ig] 2020]. Leading approaches,
et al.
2020 and batch normalization Farebrother et al. 2018 [g] et al.
namely IDAAC [Raileanu and Fergus 2021] and PPG [Cobbe al. 2021 investigate the advantages
of decoupling policy and value functions for generalization, whereas Jiang et al. [2021] propose automatic curriculum learning of levels.
different approach is add inductive bias to the neural network policy or learning algorithm.
Approaches such as Tamaa et al.] (2016), VVastelica et al. (202T). Boutilier et al.) [2020) embed a differentiable planner or learning algorithm into the neural network. Other methods (Kansky et al.) 2017. Toyer et al.] 2018 Rivlin et al.) [220)] combine learning with classical graph planning to generalize across various planning domains. These approaches require some knowledge about the problem structure (e.g., relevant planning algorithm). while our approach does not require any task-specific knowledge. Another line of work aims to learn policies or features that are invariant
across the different training tasks [Sonar al., 202 Bertran et al. 2020 Li ct al., 2021 [gl et al.,
2019 Stooke et al. 2021 Mazoure et al. 2020 and are thus more robust to sensory variations.
Using an ensemble to direct exploration to unknown areas of the state space was proposed in the model-based TEXPLORE algorithm of Hester and Stone [2013), where an ensemble of transition models was averaged to induce exploratory actions when the models differ in their prediction. Observing that exploration can help with zero-shot generalization, the model-free LEEP algorithm by Ghosh et al. [2021] is most relevant to our work. LEEP trains ensemble of policies, each on separate subset of the training environment, with loss function that encourages agreement between the ensemble members. Effectively, the KI. loss in LEEP encourages random actions when the agents of the ensemble do not agree, which is related to our method. However, random actions can be significantly less effective in exploring a domain than policy that is explicitly trained to explore, such as maximum-entrrpy policy. Consequentially, we observe that our approach leads to significantly better performance test time.
State Space Maximum Entropy Exploration Maximum entropy exploration (maxEnt, Hazan et al. 2019 MMtti et al. (20TT) is an unsupervised learning framework that trains policies that maximize the entropy of their state-visitation frequency, leading to behavior that continuously explores the environment state space. Recently, maximum entropy policies have gained attention in RL [Liu and
Abbeel. 2021b Yarats et al.) 2021 [Seo et al. [201]] Hazan et al. 2019 Mutti et al.. 2021] mainly in the context of unsupervised pre-training. In that setting, the agent is allowed to train for long period without access to environment rewards, and only during test the agent gets exposed to the reward signal and performs a limited fine-tuning adaptation learning. Importantly, these works expose the agent the same environments during pre- pre-training and test phases, with the only distinction being the lack of extrinsic reward during pre-traine gre-trannng. To the best of our knowledge, our observation that maxEnt policies generalize well in the zero-shot setting is novel.
3 Problem Setting and Background
We describe our problem setting and provide background maxEnt exploration.
3

---

Reinforcement Learning (RL)
In Reinforcement Learning an agent interacts with an unknown, stochastic environment and collects rewards. This is modeled by a Partially Observed Markov Decision Process (POMDP) (Bertsekas) 2012), which is the tuple M = (S, AS,A,O, O, P Pini. P. r),,), where S E R S| and A E R are the state and actions spaces, O is the observation space, P Pinst is an initial state distribution, P is the transition kernel, Σ is the observation function, r 5 x A + R is the reward function, and E [0. 1) the discount factor. The agent starts from initial state Su Pinit and at time performs an action on the environment that yields a reward = r(tt, 0z), and an observation Consequently, the environment transitions into the next state according to St+1 ~ P(- at). Let the history at time / be 14 {oo, (o, Yo, 0t, d1, F1 ...,, Or), the sequence of observations, actions and rewards. The agent's next action is outlined by policy т, which is a stochastic mapping from the history to an action probability T(ahh)) = P(a, alha)- In our formulation, a history- dependent policy (and not Markov policy) is required both due to partially observed states, epistemic uncertainty [Ghosh et al. 2021], and also for optimal maxEnt exploration [Mutti et al. 2022
E(st, az) O.
Zero-Shot Generalization for RL
We assume a prior distribution over POMDPs P(M), defined over some space of POMDPs. For a given POMDP, optimal policy maximizes the expected discounted return F <E 7'r((t, at)] where the expectation is taken over the policy π (hi), and the state transition probability 8 ~ Pff POMDP M. Our generalization objective in this work to maximize the discounted cumulative reward taken in expectation over the POMDP prior, also termed the population risk:
pop (T) = EMMP((M)
BR,.M
y r'r(st, at)
1=0
Secking policy that performs well in expectation over any POMDP from the prior corresponds to zero-shot generalization.
We assume access to N training POMDPs goal is to use M....., MN learn policy that performs well objective [ is to optimize the empirical risk objective:
My sampled from the prior, Mj ~ P(M). Our
A common approach
Remp (T)
1 N
ER,M,
Σxt(ss t=
= E B(MP)(
E
Σ rrss an) C
(2)
where the empirical POMDP distribution can be different from the true distribution, i.e. P(M)+ + P(M). In general, a policy that optimizes the empirical risk (Eq, ( may perform poorly on the population risk (Eq. this is known as overfitting in statistical learning theory (Shalev-Shwartz and DeniD 2014], and has been analyzed recently also for RL [Tamar et al. 2022].
Maximum Entropy Exploration
In the following provide the definitions for the state distribution and the maximum entropy exploration objective. For simplicity, we discuss MDPs-the the fully observed special case of POMDPs where O S. and (s, 5.
policy π, through its interaction with an MDP, induces a 1-step state distribution di, = p( = s|ir) over the state space S. Let (x, a) =(s 8, a =a/T) be its (-step state-action counterpart. For the infinite horizon setting, the stationary state distribution is defined as dg (s) lime-+o dr, (s), and its y-discounted version as dr
(1 7) Ee ('de,r We denote the state marginal distribution as dr,# .r (s) T E (s), which a marginalization of the 1-step state distribution over finite time T. The objective of maximum entropy exploration is given by:
H(d(-)) -
where d can be regarded as either the stationary state distribution dz Mutti and Restelli 2020 the discounted state distribution d.,,, HHaan et al., 2019) or the marginal state distribution dr. Lee et al. 2019b Mutti and Restelli 2020 In our work we focus on the finite horizon setting and

---

adapt the marginal state distribution dr in which T equals the episode horizon H, i.e. we seek to maximize the objective:
}}
Rx(T)= = EMrp =
= M~)
1 H de (s)
Σ
which yields policy that "equally" visits all states during the episode. Existing works that target maximum entropy exploration rely on estimating the density of the agent's state visitation distribution (Hazan al. 2019 Lee et al., 20196]. More recently, branch of algorithms that employ non-
parametric entropy estimation Liu and Abbeel 2021a, Mutti et al., 2021 Seo et al., 2021 has emerged, cireumventing the burden of density estimation. Here. we follow this common thread and adapt the non-parametric entropy estimation approach: we estimate the entropy using the particle- based k-nearest neighbor (k-NN estimator) [Beirlant et al. 1997 Singh et 2003 as elaborated in the next section.
4 The Generalization Ability of Maximum Entropy Exploration
In this section we present an empirical disservation-pollieee trained for maximum entropy exploration (maxEnt policy) generalize well. First, explain the training procedure of our maxEnt policy, then we show empirical results supporting this observation.
4.1
Training State Space Maximum Entropy Policy
To tackle objective 4). we estimate the entropy using the particle-based k-NN estimator [Beirlant al. 1997 Singh et al. 2003), as described here. Let
a random variable over the support C I with probability mass function Given the probability of this random variable, its entropy obtained by (p) = Exzp(log(p)] Without access to its distribution p. the entropy can be
estimated using N samples { by the k--N estimator Singh et al. 2003
N
H (p)
R
Σ
log
(1123 - x2 K-NN (12)
(5)
N
where .A-NN is the k--N sample of æ; from the set {x: J.V.
To estimate the distribution d the states S, consider each trajectory H samples of states J and take sA-NN be the k-NN of the state s, within the trajectory, proposed by previous works (APT, Liu and Abbeel 2222 la RE3, Seo et al. [2021], and APS, Liu and Abbeel (2021a)
H
H (dn,)) R
Σ H
1
log (1Is - 112)
Next, similar to previous works, since this sampled estimation of the entropy (Fq. 6 is sum of functions that operate on each state separately, it can be considered as expected reward objective (5t) with the intrinsic reward function:
H. (dH.)) ~ DH E
r (s:) == log s4 - K-NN (2)
This formulation enables us deploy any RL algorithm to approximately optimize objective (4). Specifically, in our work we use the policy gradient algorithm PPO (Schulman et al. 2017]. where at every time step the state s A-NN is chosen from previous states { of the same episode.
Another challenge stems from the computational complexity of calculating the L2 norm of the A-NN (Eq. [7) at every time step t. To improve computational efficiency, we introduce the following approximation: instead of taking the full observation as the state 8, (i.e. 64 x 64 RGB image), we sub-sample (denoted !) the observation by applying average pooling of 3 3 to produce image st of size 21 x 21, resulting in: (st) = log (1|st 1-NN.1
(8)
(2)
5

---

Maze
Jumper
Miner
00
56.
23
305
solo
5.6
1.00
5000
2/0
2000
Loves
(a) Entropy Reward
Maze
Jumper
Miner
L
Levech
Londs
(b) Extrinsic Reward
Figure 4: Generalization ability of maximum entropy vs. extrinsic reward: (a) Score of maximum entropy. (b) Score of extrinsic reward. Training for maximum entropy exhibits small generalization gap in Maze, Jumper and Miner. Average and standard deviation are obtained using 4 seeds.
We emphasize that we do not modify the termination condition of each game. However, a maxEnt policy will learn avoid termination, as this increases the sum of intrinsic rewards. In Figure we display the states visited by maxEnt policy on Maze. We also experimented with Lo as the state similarity measure instead of L2, which resulted in similar performance (see Appendix F.2).
4.2 Generalization of maxEnt Policy
The generalization gap describes the difference between the reward accu- mulated during training (T) and testing Reop of policy, where we approximate the population score by testing on a large population of tasks withheld during training. We can evaluate the generalization gap for either an extrinsic reward, or for an intrinsic reward, such as the reward that elicits maxEnt exploration (Eq. (8. In the latter, the generalization gap captures how well the agent's ' exploration strategy generalizes.
Figure 3: Example of a maxEnt trajectory on Maze. The policy vis- its every reachable state and averts termination by avoiding the goal state.
We found that agents trained for maximum entropy exploration exhibit a smaller generalization gap compared with the standard approach of training solely with extrinsic reward. Intuitively, this can be attributed to the extrinsic reward serving as 'easy" signal to learn from, and overfit to in the training environments. To assess the generalization quality of the maxEnt policy, we train agents on 200, 500, 1000 and 5000 instances of ProcGen's Maze, Jumper and Miner environments using the intrinsic reward (Eq. (). The policies are equipped with a memory unit (GRU. Cho et al. 2014) to allow learning of deterministic policies that maximize the entropy [Mutti et al. 2022
The train and test return scores are shown in Fig. 4a In all three environments, we demonstrate a small generalization gap, test performance unseen levels closely follows the performance achieved during training. When considering Maze trained on 200 levels, we observe a small generalization gap of 1.7% meaning test performance closely follows train performance. For Jumper and Miner the max EaxEnt policy exhibits a small generalization gap of 8.5% and 4.3%, respectively. In addition, verify that the train results are near optimal by comparing with hand designed approximately optimal exploration policy. For example, on Maze we use the well known maze exploring strategy wall follower, also known as the left/right-hand rule [Hendrawan] 2020 ; see Appendix B.2 for details.
Next, we evaluate the generalization gap of agents trained to maximize the extrinsic reward The results for this experiment, shown in Fig. 4b illustrate that the generalization gap for extrinsic reward
3 extensive discussion on the importance of memory for the maxEnt objective Appendix B.3 'We train for extrinsic reward using an architecture identical to that of the intrinsic reward, with the exception of the memory unit. Incorporating memory unit in this case further degrades performance (see Appendix B.3).

---

is more prominent. For comparison, when trained on 200 levels, the figure shows large generalization gap for Maze (38.8%) and Jumper (27.5%), while Miner exhibits moderate generalization gap of 13. 1 %. For an evaluation on all ProcGin games, please see Appendix B.1
5 Explore to Generalize (ExpGen)
Our main insight is that, given the generalization property of the entropy maximization policy established above, an agent can apply this behavior in test MDP and expect effective exploration at test time. In the following, we pair this insight with the epistemic POMDP idea, and propose to play the exploration policy when the agent faces epistemic uncertainty, hopefully driving the agent different state where the reward-secking policy is more certain. This can be seen as an adaptation of the seminal explicit explore or exploit idea [Kearns and Singh 2002 to the setting of ZSG-RL.
5.1 Algorithm
Our framework comprises two parts: an entropy maximizing network and an ensemble of networks that maximize extrinsic reward to evaluate epistemic uncertainty. The first step entails training network equipped with memory unit to obtain a maxEnt policy that maximizes entropy. as
described in section 4.1 Next, we train an ensemble of memory-less policy networks { mm to maximize extrinsic reward. Following Ghosh et al. [2021], we shall use the ensemble to assess epistemic uncertainty. Different from Ghosh et al. [202]] however, we do not change the RL
loss function, and use an off-the-shelf RL algorithm (such PPO [Schulman et al. 2017 or IDAAC (Stooke et al. 202T). )
At test time, we couple these two components into
Algorithm Explore Generalize (ExpGien)
a combined agent π (de-
1: Input:
ensemble size m, initial state 80 = EXVIRONMENT.resel
tailed pseudo-code in Al- gorithm[]. I We consider do- mains with finite action
2:
Buy
4: Train maxEnt policy TयE using intrinsic reward (Eq:
space, and say that the pol- icy of is certain at state 8 if its action ai x (a]s) is) is in consensus with the ensem- ble: O; aj for the ma- jority of & out of m, where k is hyperparameter of our algorithm. When the networks { are not in consensus, the agent æ takes sequence of nTw tions from the entropy max- imization policy THe which encourages exploratory be- havior.
Train m policies n.. m. T.. using extrinsie reward Tezt.
6: for t = to do
a; ~ (-(ss
8:
nen + Rxu -1 if 0; E Consensus(ay, iie (.....m) and now < 0 then =
10:
11:
else
12:
az ~ (-(()))))
13:
a₁ = Hu
14:
nh ~ Geom(a)
15:
end if
16:
S1+1
ENVIRONMENT.STEP(()
17:
end for
Agent mela-stability Switching between two policies may result in case where the agent repeat- edly toggles between two states-if, say, the maxEnt policy takes the agent from state 8j to a state 82. where the ensemble agrees on an action that again moves to state s1. To avoid such "meta-stable" behavior, we randomly choose the number of maxEnt steps nTy from a Geometric distribution, MTZ Geom(a).
Experiments
We evaluate our algorithm on the ProcGen benchmark, which employs discrete 15- 15-dimensional action space and generates RGB observations of size x x 3. Our experimental setup follows ProcGen's 'easy' configuration, wherein agents are trained on 200 levels for 25M steps and subse- quently tested on random levels [Cobbe et al. 22020.. All agents are implemented using the IMPALA
convolutional architecture [Espcholt et al. 2018 and trained using PPO [Schulman et al. 2017) or
7

---

0=5
y= 0.9
y= 0.99
8-
Test Score
6-
6-
4-
2
2
2
0
Step
1e7
Step 1e7
Step 1e7
Figure 5: Test performance of PPO trained using the reward Itotal that combines intrinsic and extrinsic rewards, weighted by 3 (Eq. 9 Each figure details the results for different values of discount factor y. All networks are randomly initialized and trained on 200 maze levels, and their mean is computed over runs with different seeds. The figures show an improvement over the PPO baseline for y = 0.5. In all cases, ExpCie outperforms the combined reward agent.
IDAAC [Raleaau and Fergus 2021]. For the maximum entropy agent TH. we incorporate a single GRU [hoo et al.[22014 at the final embedding of the IMPALA convolutional architecture. For all games, we use the same parameter a = 0.5 of the Geometric distribution and form an ensemble of 10 networks. For further information regarding our experimental setup and specific hyperparameters, please refer to Appendix C
6.1 Generalization Performance
We compare our algorithm to six leading algorithms: vanilla PPO [Schulman et al. 2017 PLR Jiang et al. 2021] that utilizes automatic carriccuuumbbsseee learning, UCB-DrAC [Ralleanu et al. 2021 which incorporates data augmentation to learn policies invariant to different input transformations, PPG [Cobbe et al., 2021], which decouples the optimization of policy and value function during learning, and IDAAC [Raileanu and Fergus) 2021]. the previous state-of-the-art algorithm on ProcGen that decouples policy learning from value function learning and employs adversarial loss to enforce invariance to spurious features. Lastly, we evaluate our algorithm against LEEP [Ghosh et al.] [021], the only algorithm that, to our knowledge, managed to improve upon the performance of vanilla PPO on Maze and Heist. The evaluation matches the train and test setting detailed by the contending algorithms and their performance is provided reported by their authors. For evaluating LEEP and IDAAC, we use the original implementation provided by the authors.
Tables 2 and 1 show the train and test scores, respectively, for all ProcGen games. The tables show that ExpGen combined with PPO achieves notable gain over the baselines on Maze, Heist and Jumper, while on other games, invariance- based approaches perform better (for example, IDAAC leads on BigFish, Plunder and Climber, whereas PPG leads on CaveFlyer, and UCB-DrAC leads on Dodgeball). These results correspond to our observation that for some domains, invariance cannot be used to completely resolve epistemic uncertainty. We emphasize that ExpGen substantially outperforms LEEP all games, showing that improved exploration test time is significant. In Appendix we compare ExpGen with LEEP trained for 50M environment steps, showing a similar trend. When combining ExpGen with the leading invariance-baset approach, IDAAC, we establish that ExpGen is in-fact complementary to the advantage of such algorithms, setting a new state-of- the-art performance in ProcGen. A notable exception is Dodgeball, where all current methods still fail.
Figures and 7 show aggregate statistics of ExpGen, PPO, PLR, UCB-DrAC, PPG and IDAAC for all games affirming the dominance of IXpGen+IDAAC as the sfate-of-the-art The results are obtained using 10 runs per game, with scores normalized as in Appendix C.I The shaded regions indicate 95% Confidence Intervals (CIs) and are estimated using the percentile stratified bootstrap with 2, 000 (Fig. (6) and 50, 000 (Fig.) bootstrap re-samples. Fig. 6 (Left) compares algorithm score-distribution, illustrating the advantage of the proposed approach across all games. Fig. (
Foor LEEP and IDAAC, we followed the prescribed hyperparameter values of the papers' authors (we directly corresponded with them). For some domains, we could not reproduce the exact results, and those cases, we used their reported scores, giving them an advantage.
"See rliable (Agaraal et al. 2021 for additional details the various performance measures and protocols.
8

---

Algorithm X EXpGen-IDAAC
Algorithm
1.20
IDAAC
score
0.75
EXpGen-IDAAC
PPG
EXpGen-IDAAC
UCB-DrAC
EXpGen-IDAAC
PPO
IDAAC
PPG
0.25
IDAAC
UCB-DrAC
IDAAC
PPO
0.7
0.8
0.9
1.0
Min-Max Nammlized Test Score (r)
P(X Y)
Figure 6: Comparison across all ProcGen games, with 95% bootstrap CIs highlighted in color. Left. Score distributions of ExpGen, PPO, PLR. UCB-DrAC, PPG and IDAAC. Right. Shows in each row, the probability of algorithm X outperforming algorithm Y. The comparison illustrates the superiority of ExpGen over the leading contender IDAAC with probability 0.6, as well as over other methods with even higher probability.
shows the probability of improvement of algorithm X against algorithm Y. The first row (ExpGen vs. IDAAC) demonstrates that the proposed approach surpasses IDAAC with probability 0.6 and subsequent rows emphasize the superiority of ExpGen over contending methods at even higher probability. This is because ExpGen improves upon IDAAC in several key challenging tasks and is on-par in the rest. Fig. provides aggregate metrics of mean, median and IQM scores and optimality gap (as mean) for all ProcGen games. The figure shows that ExpGen outperforms the contending methods in all measures.
Ablation Study One may wonder if the ensemble in ExpGen is necessary, or whether the observa- tion that the maxEnt policy generalizes well can be exploited using single policy. We investigate the effect of combining the intrinsic and extrinsic rewards, Yy and Tezts respectively, into single reward as a weighted sum:
Ttotal =Br (1 B)Text.
(9)
and train for 3 steps for different values of discount factor y. We obtain the best test score for y 0.5 and 0.1, illustrating an improvement compared with the PPO baseline. When comparing with ExpGen, the combined reward (Eq. 9 exhibits inferior performance with slightly higher variance. Appendix C.2 and we also provide an ablation study of ensemble size and draw comparisons to variants of our algorithm.
on Maze. Figure shows the train and test scores over 50M
Game
PPO
PLR
UCB-DrAC
PPG
IDAAC
LEEP
ExpGen (PPO)
ExpGen (IDAAC)
Bighish
2.9 1.1
10.9 + 2.8
9.2 2.0
11.2 + 1.4
18.5
0.9
6.0 0.5
18.5
StarPilot
24.9 1.0
27.9 4.4
30.0
47.2
3.3
3.2
0.9
39.8 29.
FruitBot
26.2
28.0
27.6
27.8
0.5
16.4
2..4 0.6
BossFight
7.4
8.9
+.6
1.2
9.8
0.5 0.3
+.2
9.8 0.5
Ninja
6.2
7.2
6.6
6.6
6.8
4.4
6.6
6.6 1.6.0.3
Plunder
7.6
8.7
8.3
14.3220
2..3
4.4
5.3
23.6
CaveFlyer
0.5
0.5
5.0
7.0
5.0
4.2
5.7
0.7
CoinRun
0.2
0.5
8.2
8.9
9.4
7.3
8.1
0.3
Jumper
5.8
5.8
6.3
5.9+0.11
6.3+0.2
5.4
6.7+
6.8
Chaser
3.1
6.9
6.3
9.+
6.8
3.0
3.6+166
7.1 1.4
Climber
0.4
6.3
0.6
28400.4
8.340.4
2.6
5.9
9.5
Dodgeball
2.2
1.8
4.200.9
2.3
3.2+0.3
1.9
2.9
2.8+0.2
Heist
0.5
0.5
3.5+0.4
2.8
3.5+
4.5
7.4
7.2
Leaper
4.9
6.2
4.8
8.5
7.7
4.4
4.0
Maze
5.6
5.5
6.3
5.1
5.6
0.2
8.3 0.3
7.8 0.2
Miner
7.8
9.6 0.6
9.2 0.6
7.4
9.5
1.100.1
8.0 0.7
9.8 1.8.000 0.3
Table 1: Test score of ProcGen games trained on 200 levels for 25M environment steps. We compare our algorithm to PPO, PLR, UCB-De -CB-DrAC. PPG. IDAAC and LEEP. The mean and standard deviation are computed over 10 runs with different seeds.

---

Game
PPO
PLR
UCB-DrAC
PPG
IDAAC
LEEP
ExpGen (PPO)
ExpGen
(IDAAC)
BigFish
8.9
7.8
12
21.8
0.9
7.0
21.5 2.3
StarPilot
29.0 1.1
2.6 0.3
33.1 + 1.3
49.6 2.1
38.6 2.2
5.3
34.3 1.6
40.0
FruitBot
28.8 + 0.6
15.9
29.3 0.5
+.5
29.1
1..4
28.9 0.6
29.5
BossFiggt
8.0
8.7
8.1
11.1
1.44
0.3
7.9
9.9
Ninja
7.3
0.5
8.4004
8.2
8.3
4.6
8.5
7.9
Plunder
9.4++1.7
4.1+13
10.2+188
16.4+1199
1.6
5.8
2.71
CaveFlyer
0.7
0.1
5.9
9.5
6.2.1066
4.9
6.8
5.5
CoinRun
9.4 0.3
5.4
9.4022
9.9
9.8
6.7
9.8
9.1
Jumper
8.6
0.5
8.1
8.7
8.7
5.7
7.9
8.1
Chaser
3.7
0.7
7.0
10.7
7.5
2.6
4.7
Climber
6.9-
0.8
8.+
1.2
1..2
3.5+
7.7+0.4
11.4
Dodgehall
6.6
2.0
7.3+0.8
5.5
4.9
3.3+
5.8
5.3 0.5
Heist
6.1
1.2004
6.2006
7.40.44
4.5
7.1
9.10.
7.0
Leaper
5.4
0.4
5.0
9.111
8.3
4.3
43120
8.1
Maze
9.1++.2
4.1
8.5+
9.0
6.4
9.4
9.6
7.4+0.4
Miner
0.3
9.7
1.3
9.0
11.9
Table 2: Train score of ProcGen games trained on 200 levels for 25. environment steps. We compare algorithm to PPO, PLR, UCB-DrAC, PPG, IDAAC and LEEP. The mean and standard deviation are computed over runs with different seeds.
Median
IQM
Mean
Optimality Gap
EXpGen-IDAAC
IDAAC
PPG
PLR
UCB-DrAC
PPO
0.3 0.4 0.5
0.6
0.3 0.4 0.5 0.6
0.4 0.5 0.6
0.4
0.5
0.6
0.7
Min-Max Normalized Test Score
Figure 7: Aggregate metrics for all ProcGen games: mean, median and IQM scores (higher is better) and optimality gap (lower better), with 95% Cls highlighted in color. ExpGen outperforms the contending methods in all measures.
7 Discussion and Limitations
We observed that policies trained to explore, using maximum entropy RL, exhibited generalization of their exploration behavior in the zero-shot RI. setting. Based on this insight, we proposed ExpGen- ZSG-RL algorithm that takes maxEnt exploration step whenever ensemble of policies trained for reward maximization does not agrec on the current action. We demonstrated that this simple approach performs well all ZSG-RL domains of the ProcGen benchmark.
One burning question is why does maxEnt exploration generalize so well? An intuitive argument is that the maxEnt policy in an MDP is invariant to the reward. Thus, if for every training MDP there are many different rewards, each prescribing different behavior, the maxEnt policy has to be invariant to this variability. In other words, the maxEnt policy contains no information about the rewards the data, and generalization well known to be bounded by mutual information between the policy and the training data [Bassily et al.) 2018). Perhaps an even more interesting question is whether the maxEnt policy also less sensitive variations the dynamics the MDPs. We leave this open thcoretical problem.
Another consideration is safety. some domains, wrong action can lead disaster, and in such cases, exploration at test time should be hedged. One possibility is to add to ExpGen's 's policy ensemble an ensemble of advantage functions, and use weigh the action agreement [Rotman et al.) 2020). Intuitively, the ensemble should agree that unsafe actions have low advantage, and not select them at test time.
Finally, we point out that while our work made significant progress on generalization in several ProcGien games, the performance Dodgeball remains low for all methods
aware of. An
interesting question is whether performance on Dodgeball can be improved by combining invariance- time, or whether Dodgeball represents
based techniques (other than IDAAC) with exploration a different class of problems that requires a completely different approach.
10

---

The research of DS was Funded by the European Union (ERC, A-B-C-Deep. 101039436). The research of EZ and AT was Funded by the European Union (ERC, Bayes-RL. 101041250). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union the European Research Council Executive Agency (ERCEA). Neither the European Union nor the granting authority can be held responsible for them. DS also acknowledges the support of the Schmidt Career Advancement Chair Al.
Acknowledgments
References
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304-299 29320, 2021.
Raef Bassily, Shay Moran, Ido Nachum, Jonathan Shafer, and Amir Yehudayoff. Learners that use little information. In Algorithmic Learning Theory, pages 25-55, PMLR, 2018.
Jan Beirlant, Edward Dudewicz, László Györfi, Edward C Van der Meulen, al. Nonparametric entropy estimation: An overview. International Journal of Mathematical and Statistical Sciences,
6(1):17-39, 1997.
Martin Bertran, Natalia Martinez, Mariano Phielipp, and Guillermo Sapiro. Instance-based gen- eralization in reinforcement learning. Advances in Neural Information Processing Systems, 33: 11333-11344, 2020.
Dimitri Bertsekas. Dynamic programming and optimal control: Volume I, volume 1. Athena scientific, 2012.
Craig Boutilier, Chih-wei Hsu, Branislay Kveton, Martin Mladenov, Csaba Szepesvari, and Manzil Zaheer. Differentiable meta-learning of bandit policies. Advances Neural Information Processing Systems, 33, 2020.
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint urXiv: 1259, 2014.
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In International Conference on Machine Learning, pages 1282-1289. PMLR, 2019.
Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, pages PMLR, 2020.
Karl W Cobbe, Jacob Hilton, Oleg Klimov, and John Schulman. Phasic policy gradient. In Interna- tional Conference on Machine Learning, pages 2020-2027. PMLR, 2021.
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International conference on machine learning, pages 1407-1416. PMLR, 2018
Jesse Farebrother, Marlos C Machado, and Michael Bowling, Generalization and regularization in dqn. arXiv preprint arXiv:1810.0011 2018.
Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang, Ryan Adams, and Sergey Levine. Why generalization in rl is difficult: Epistemic pomdps and implicit partial observability, Advances in Neural Information Processing Systems, 34:25502-25515, 2021.
Elad Hazan, Sham Kakade, Karan Singh, and Abby Van Soest. Provably efficient maximum entropy exploration. International Conference on Machine Learning, pages 2681-2691. PMLR, 2019.
YF Hendrawan. Comparison of hand follower and dead-end filler algorithm solving perfect mazes. In Journal of Physics: Conference Series, volume 1569, page 022059. IOP Publishing. 2020.

---

Todd Hester and Peter Stone. Texplore: real-time sample-efficient reinforcement learning for robots. Machine learning. 90:385 -289, 2013.
Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, and Katja Hofmann. Generalization in reinforcement learning with selective noise injection and information bottleneck. Advances in neural information processing systems, 32, 2019.
Maximilian Igl, Gregory Farquhar, Jelena Luketina, Wendelin Bochmer, and Shimon Whiteson. The impact of non-stationarity on generalisation in deep reinforcement learning. arXiv preprint arXiv:2006.000 05826, 2020.
Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. In International Conference on Machine Learning, pages 4940-4950. PMLR, 2021.
Ken Kansky, Tom Silver, David A Mély, Mohamed Eldawy, Miguel L.ázaro-Gredilla Xinghua Lou, Nimrod Dorfman, Szymon Sidor, Scott Phoenix, and Dileep George. Schema networks: Zero-shot transfer with a generative causal model of intuitive physics. In International conference on machine learning, pages 189--1818. PMLR, 2017.
Michael Kearns and Satinder Singh. Npar-optimal reinforcement learning in polynomial time.
Machine learning. 49(2):209 232, 2002.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 1412.6900, 2014.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of generalisation in 2021.
deep reinforcement learning. arXiv preprint
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprini 23X1v:2004.13649, 2020.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization deep reinforcement learning. arXiv preprint arXiv:1910.053366 2019a.
Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric Xing, Sergey Levine, and Ruslan Salakhutdinov. Efficient exploration via state marginal matching. arXiv preprint 1rXiv:/906.0530 2019b.
Bonnie Li, Vincent François-Laved Thang Doan, and Joelle Pincau. Domain adversarial reinforcement learning. arXiv preprint arXiv:2102.07097, 2021.
Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. In International Conference on Machine Learning, pages 6736-6747. PMLR, 2021a.
Hao Liu and Pieter Abbeel. Behavior from the void: Unsupervised active pre- preitraining. Advances in Neural Information Processing Systems, 34: 14:19499-18473, 2021b.
Bogdan Mazoure, Remi Tachet des Combes, Thang Long Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. Advances in Neural Information Processing 2020.
Systems,
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 518 5333333, 2015.
Mirco Mutti and Marcello Restelli. An intrinsically- mottaatee approach for learning highly exploring and fast mixing policies. In Proceedings of the AAAI Conference on Artificial Intelligence. volume 34, pages 5232-5233 2020.
Mirco Mutti, Lorenzo Pratissoli, and Marcello Restelli. Task-agnostic exploration via policy gradient of a non- nor-mmarttr state entropy estimate. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9028-9036. 2021.
Mirco Mutti, Riceardo De Santi, and Marcello Restelli. The importance of non-markovianity in maximum state entropy exploration. arXiv preprint arXiv:2202.03066, 2022.
12

---

Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization in reinforcement learning. International Conference on Machine Learning, pages 8787-8798. PMLR, 2021.
Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in reinforcement learning. Advances in Neural Information Processing Systems, 34:5402-540 5515, 2021.
Or Rivlin, Tamir Hazan, and Erez Karpas. Generalized planning with deep reinforcement learning. arXiv preprint 2xXiv:0005002303 2020.
Noga H Rotman, Michael Schapira, and Aviv Tamar. Online safety assurance for beaning--uggeenne systems. In Proceedings of the 19th ACM Workshop on Hot Topics in Networks, pages 88-95, 2020.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604-609, 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: (7XX:v1707.063477 2017.
Younggyo Sco, Lili Chen, Jinwoo Shin, Honglak Lee, Pieter Abbeel, and Kimin Lee. State entropy maximization with random encoders for efficient exploration. In International Conference on Machine Learning, pages 9443-9454. PMLR, 2021.
Shai Shalev-Shwartz and Shai Ben-David. Understanding machine learning: From theory to algorithms. Cambridge university press, 2014.
Harshinder Singh, Neeraj Misra, Vladimir Hnizdo, Adam Fedorowicz, and Eugene Demchuk. Nearest neighbor estimates of entropy. American journal of mathematical and management sciences, 23 (3-4):301-321, 2003.
Anoopkumar Sonar. Vincent Pacelli, and Anirudha Majumdar. Invariant policy optimization: Towards stronger generalization in reinforcement learning. In Learning for Dynamics and Control, pages 21-33. PMLR, 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning, pages 9870-9879, PMLR, 2021.
Aviv Tamar, Yi Wu, Garrett Thomas, Sergey Levine, and Pieter Abbeel. Value iteration networks. Advances in neural information processing systems, 29, 2016
Aviy Tamar, Daniel Soudry, and Ev Zisselman. Regularization guarantees generalization in bayesian reinforcement learning through algorithmic stability. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8423-8431, 2022.
Sam Toyer, Felipe Trevizan, Sylvie Thiébaux, and Lexing Xie. Action schema networks: Generalised policies with deep learning. In Thirty-Second AAAI Conference on Artificial Intelligence, 2018.
Oriol Vinyals, Igor Babuschkin, Wojciech M Czamecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung. David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 575(7782):350-3 2019.
Marin Vlastelica, Michal Rolínek, and Georg Martius. Neuro-algorithmic policies enable fast combinatorial generalization. arXiv preprint arXiv:2102.07456, 2021.
Kaixin Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforcement learning with mixture regularization. Advances in Neural Information Processing Systems, 33: 7968-7978, 2020.
Peter R Wurman, Samuel Barrett, Kenta Kawamoto, James MacGlashan, Kaushik Subramanian, Thomas J Walsh, Roberto Capobianco, Alisa Devlic, Franziska Eckert, Florian Fuchs, et al. Outracing champion gran turismo drivers with deep reinforcement learning. Nature, 602(7896): 223-228, 2022.
13

---

Huaxiu Yao. Linjun Zhang, and Chelsea Finn. Meta-learning with fewer tasks through task interpola- tion. arXiv preprint arXiv:2106.0269) 02695, 2021.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In International Conference on Machine Learning, pages 11920- 11931. PMLR, 2021.
Chang Ye, Ahmed Khalifa, Philip Bontrager, and Julian Togelius. Rotation, translation, and cropping for zero-shot generalization. In 2020 IEEE Conference Games (CoG), pages 57-64. IEEE, 2020.
14

---

A Hidden Maze Experiment
In this section empirically validate thought experiment described in the Introduction: train a recurrent policy on hidden mazes using 128 training levels (with fixed environment colors). Tigure[8adepiin the observations of an agent along its trajectory; the agent sees only its own location (green spot), whereas the entire maze layout, corridors and goal, are hidden. In Fig. 8b we visualize the full, unobserved, state of the agent. The train and test results are shown in Fig indicating severe overfitting to the training levels: test performance failed to improve beyond the random initial policy during training. Indeed, the recurrent policy memorizes the agent's training trajectories instead of learning generalized behavior. Hiding the maze's ' goal and corridors leads to agent-behavior behavior that is the most invariant instance specific observations the observations include the minimum information, such that the agent can still learn to solve the maze.
2
12
(a) Agent observations along a trajectory the Hidden Maze task. The bottom-right frame 12) shows agent eventually revealing the goal.
1-9
1-11
1-
L
LI
(b) corresponding unobserved true states agent. Unexplored regions the maze appear faded sections.
Figure 8: Hidden Maze experiment where the agent only observes its own location (green spot). Both the goal (purple spot) and corridors are not observable (appear as walls).
This experiment demonstrates why methods based on observation invariance (e.g., IDAAC Raileanu and Fergus (202T) do not improve performance on Maze-like tasks, despite significantly improving performance on games such BigFish and Plunder, where invariance colors helps to generalize.
15

---

Hidden Maze Experiment
10
Score
6-
2
4
Step
1e7
Figure 9: PPO performance on the hidden maze task, indicating severe overfitting. Train (Orange)
and Test (blue) performance are displayed for 10 seeds, alongside their means (highlighted in bold).
Maximum Entropy Policy
This section elaborates on the implementation details of the maxEnt oracle and provides a performance evaluation of the maxEnt policy.
B.1 Generalization Gap of maxEnt vs PPO across all ProcGen Environments
Recall from Section 4.2 that the (normalized) generalization gap is described by
(Remp (T) Roop (mm)//mm
where Remp(1) (T) Rtrai (i)) and Pppp(T) (TT) Reest (T). Fig 10 shows the generalization ability of
the maxEnt exploration policy compared to PPO, obtained from training on 200 training levels.
Generalization Gap (%) of Maximum Entropy Policy (lower is better)
maxEnt Gap [%]
PPO Gap [%)
60
30
Bigfish
StarPittt
FruiBot
BossFight
Nisja
Plunder
Caveflyer
ColnRun
jumper
Chasee
Climber
Dodgeball
Heier
Lepper
Maze
Miner
Figure 10: The normalized generalization gap [%] of maxEnt and PPO for all ProcGen games, trained on 200 training levels and averaged over 4 seeds (lower is better).
The figure demonstrates that the max exploration policies transfer better in zero-shot generalization (achieve smaller generalization gap) across all ProcGen games apart from Ninja. This holds true
16

---

even in environments where the ExpGen algorithm is on par but does not exceed the baseline, pointing to the importance of exploratory behavior in some environments, but not in others.
B.2 Computing the maxEnt Oracle
The maxEnt oracle prescribes the maximum intrinsic return achievable per environment instance. To simplify the computation of the maxEnt oracle score, in this section, we evaluate the maxEnt score using the Lu instead of the L2 norm and use the first nearest neighbor (k = 1). In Maze, we implement the oracle using the right-hand hand rule [Hendrawan 2020] for maze exploration: If there is no wall and no goal on the right, the agent turns right, otherwise, it continues straight. If there is wall (or goal) ahead, it goes left (see Figure 33. In Jumper, we first down-sample (by average pooling) from 64 x 64 to 21 x 21 pixels to form cells of 3 x 3, and the oracle score is the number of background cells (cells that the agent can visit). Here we assume that the agent has size of 3 x 3 pixels. In Miner, the "easy" environment is partitioned into 10 x 10 cells. The maximum entropy the number of cells that contain "dirt" (i.e., cells that the agent can excavate in order to make traversable).
Maze
Jumper
Miner
Spope Oredte
Score cs. Oredle
6.6
0.6
0.6
Scuress,
0.4
0.4
0.4
Train with GRU
6.2
Test with GRU
0.2
0.2
Train withour GaU
That whthou GRii
5.0
0.0
300
0.0
520
Levele
Levukk
Maze
Jumper
Miner
1.0
1.0
1.0
6.8
0.8-
0.8-
6.6
0.6
0.G.
Supcess
0.6
0.4
0.4
0.2
0.2
0.2
200
5:0
1000
5030
0.0
200
502
1000
5000
0.0
205
510
1000
Soco
Lovela
Level
Leves
(a) Entropy Reward
Maze
Jumper
Miner
20
12
Siore
200
500
1002
5000
200
100
2000
5000
200
520
2000
5000
Levett
Level
Levex
(b) Extrinsic Reward
Figure/ 11: Generalization ability of maximum entropy and extrinsic reward policy: (a.top row) Score of maximum entropy policy, normalized by the oracle score. (a.bottom row) Success Rate of maximum entropy policy. (b) Score of extrinsic reward policy. Training for maximum entropy exhibits a small generalization gap in Maze, Jumper and Miner. Average and standard deviation are obtained using 4 seeds.
17

---

In Fig. Iaa the top row describes the intrinsic return of the maxEnt policy, normalized by the oracle's return (described above). We draw a comparison between agents with memory (GRU) and without. For the Maze environment (top-right), the agent achieves over 90% of the oracle's performance when employing memory unit (GRU). For the Jumper and Miner environments, the agents approach 60% and 80% of the oracle's & score, respectively.
Next, the bottom row of Fig. IIa details the success rate as the ratio of instances in which the agent successfully reached the oracle's 's score (meaning that the entire traversable region has been explored). For the Maze environment, the agent achieves a success-rate of 70%, whereas for the Jumper and Miner the agent fails to meet the oracle's 's score. This indicates that the oracle, as per our implementation for Jumper and Miner, captures states that are in-effect unreachable, and thus the agent is unable to match the oracle's score.
B.3 The Importance of a Memory Unit for maxEnt
We evaluate the importance of memory for maxEnt: In Fig. Ma train maxEnt agent with and without memory (GRU). The results indicate that for Maze, the GRU vital in order to maximize performance for all various sizes of training set. In Jumper, advantage when training on 200 levels, however the benefit diminishes when additional training levels available. For Miner, there appears be no significant advantage for incorporating memory with 200 training levels. However, the benefit of a GRU becomes noticeable once more training levels are available (beyond 500).
that the GRU provides an
When looking at the extrinsic reward (Fig. [lb), we see an interesting effect after introducing a GRU. Maze and Jumper suffer degradation in performance with 200 levels (indicating overfit), while Miner appears to be unaffected.
In summary, we empirically show that memory is beneficial for the maxEnt policy on Maze, Jumper and Miner. Interestingly, we demonstrate that the introduction of memory for training maximize extrinsic reward causes the agent to overfit in Maze and Jumper with 200 training levels.
C Experimental Setup
This section describes the constants and hyperparameters used part of the evaluation of our algorithm.
C.1 Normalization Constants
In Figures 2 aand 12 we compare the performance of the various algorithms. The results are normalized in accordance with [Cobbe et al., 2020), which defines the normalized return as
Rnorm (R Rmmin) (Bmar Rmmin.
The Rmin and Rmar of each environment detailed in Table Note that the performance of PPO on Heist (see Fig. is lower than Rmin (the trivial performance), indicating severe overfitting.
C.2 Hyperparameters
As described in Section[5.1 the hyperparameters of our algorithm are the number of agents that form the ensemble, which i agents are required to be agreement for the ensemble to achieve consensus on its action, and a as the parameter of MTw Geom((:) that represents the number of maxEnt steps taken when the ensemble fails to reach a consensus. An additional hyperparameter is the neighborhood size ANN of the k-NN estimator (Section 4.1b. used to determine the reward of the maxEnt policy.
We conducted a hyperparameter search over the ensemble size E {4, 6, 8, 10) for different values of ensemble agreement & E {2, 4, 6. 8), and values of (0.2, 0.8}. We found that value of a 0.5, and an ensemble size of m 10 produce the best results for all games, whereas the value of varies from game to game, as detailed in table For example, table[5]shoss the results for varying values of m and & for the Maze environment. Throughout our experiments, we train our networks using the Adam optimizer [Kingma and Bal 2014]. For the PPO hyperparameters we use the hyperparameters found in [Cobbe al. 2020 as detailed in Table 6
18

---

Hard
Easy
Environment
Rmin
Rmaa
Rmin
Rmaz
CoinRun
5
10
5
10
StarPilot
1.5
35
2.5
64
CaveFlyer
13.4
3.5
12
Dodgeball
1.5
19
1.5
19
FruitBot
-.5
27.2
-1.5
32.4
Chaser
.5
14.2
.5
13
Miner
1.5
20
1.5
13
Jumper
10
3
10
Leaper
1.5
10
3
10
Maze
10
5
10
BigFish
40
40
Heist
2
10
3,5
10
Climber
12.6
2
12.6
Plunder
3
30
4.5
30
Ninja
10
3.5
10
BossFight .5
13
.5
13
Table 3: Normalization Constants.
Game
Maze
Jumper
Miner
Heist
BigFish
StarPilot
FruitBot
BossFight
Plunder
CaveFlyer
2
CoinRun
Chaser
2
Climber
Dodgeball
2
Leaper
Ninja
2
Table Consensus size & as hyperparameter for each game.
Ensemble size ExpGen
4
10
8.02 + 0.06
8.15 - 0.19
8.00 0.12
8.22 + 0.11
Table 5: Ablation study of ensemble size and its effect the test score. Each network in the ensemble trained on 200 instances of Maze. The results show improved performance for large ensemble size. The mean and standard deviation computed using 10 runs with different seeds.
Table 7 shows an evaluation of the maxEnt gap and ExpGen using different neighhor size kyy on the Maze environment. We found that the best performance obtained for KNN 21, 5). Thus, we choose ĀNN 2, the second nearest neighbor, for all games.
D Results for all ProcGen Games
Figure 12 details the normalized test performance for ProcGen games. Normalization per- formed according to [obbee et al., 2020 as described in Appendix C.I The figure demonstrates
19

---

Parameter
Value
X
.999
# timesteps per rollout Epochs per rollout
512
# minibatches per epoch
Entropy bonus (k#)
.01
PPO clip range
.2
Reward Normalization?
Yes
Learning rate
5e-4
# workers
# environments per worker
32
Total timesteps
25M
GRU?
Only for maxEnt No
Frame Stack?
Table 6: PPO Hyperparameters.
Neighbor
Maze
Environment maxEnt Gap [%]
Size KNN
maxEnt (Train)
maxEnt (Test)
ExpGen (Test)
17.5 + 0.5
15.3 = 1.4
12.7%
7.9 0.2
2
33.9 + 0.2
31.3 + 1.5
7.7%
8.3 0.2
3
50.2 + 2.0
42.5 + 2.8
15.3%
8.2 0.2
4
62.9
52.3 3.0
16.8%
8.1
70.6
57.8 4.6
15.5%
8.2
80.0 3.2
65.3
18.4%
8.1
7
87.2 1 1.4
71.1 1.1
18.5%
8.0
8
03.8 1.0
75.3 2.3
19.7%
7.8 0.2
96.6 4.8
78.0
19.3%
8.1 0.2
10
102.2 = 1.8
89.1
12.8%
7.9 1 0.1
Table 7: Hyperparameter search for neighborhood size kNx for the Maze environment. The table presents the maxEnt gap and the performance ExpGen for varying values kNN. The mean and variance are computed for seeds.
that ExpGen establishes state-of -oie---ftthhttte art results on several challenging games and achieves on-par performance with the leading approach on the remaining games.
E
Results after convergence
Tables g and | detail the train and test performance of ExpGen, LEEP and PPO, when trained for 50M environment steps. Table shows that ExpGen surpasses LEEP and PPO in most games.
Game
PPO
LEEP
ExpGen
Maze
9.45 + 0.21
9.84 + 0.05
9.71 0.11
Heist
7.97 0.56
6.86 0.68
9.62 0.08
Jumper
8.62 0.08
6.1
8.00
Miner
12.86 = 0.06
1.9
12.53 0.122
BigFish
14.24 + 3.36
0.36
5.95 0.32
Climber
8.76 0.41
4.6
8.89 +.29
Dodgeball
8.88 0.38
6.22
8.73 0.40
Plunder
9.64 + 1.85
5.1
8.06 + 0.98
Ninja
9.10 0.32
5.2 0.3
8.90 + 0.30
CaveFilyr
8.98 + 0.59
6.4
8.75 + 0.59
Table 8: Train score of ProcGen environments trained on 200 instances for 50M environment steps. We compare our algorithm to the baselines LEEP and PPO. The mean and standard deviation are computed over 8 runs with different seeds.
20

---

Normalized test performance on all ProcGen games
1.2
PPO
EXpGen+IDAAC
0.8
0.4
0.0
-02
CainRun
StarPilot
CaveFiyer
Dodgeball
oruitBot
Chaser
Miner
Jumper
Leaper
Maze
Bigfish
Heist
Climber
Plunder
Ninja
BossFight
Figure 12: Normalized test Performance for PPO, IDAAC, and ExpGen+IDAAC, on all ProcGen games. ExpGen achieves state-of the art performance test levels of Maze, Heist, and Jumper and on-par performance in the remaining games.
Game
PPО
LEEP
ExpGen
Maze
5.78 0.39
6.78
0.13
Heist
2.54 0.45
0.57
6.21
Jumper
5.78 0.28
0.4
6.14
Miner
8.76 + 0.33
0.1
9.38
BigFish
3.82 1.98
5.5
0.69
Climber
6.14 0.50
2.6
6.29
Dodgeball
3.71 0.55
0.47
3.84
Plunder
7.78 + 1.74
4.2
6.91 + 1.00
Ninja
6.94 0.30
4.9 0.8
6.75 6.75
CaveFlyer
6.19 + 0.66
2.6 + 0.2
6.36 + 0.49
Table 9: Test score of ProeGen environments trained on 200 instances for 50M environment steps.
We compare our algorithm to the baselines LEEP and PPO. The mean and standard deviation computed over 8 runs with different seeds.
F
Ablation Study
In the following sections, we provide ablation studies of an ExpGen variant that combines random actions and an evaluation of Ly state similarity measure.
F.1 Ensemble Combined with Random Actions
compare the proposed approach variant ExpGen denoted by Ensemble+random where we train an ensemble and at test time select random action if the ensemble networks fail to reach consensus. The results are shown in Table indicating that selecting the maximum entropy policy upon ensemble disagreements yields superior results.
21

---

Algorithm
Train
Test
ExpCien
9.6 0.2
8.1
Ensemble + random
9.3 0.1
0.2
LEEP
9.4
6.6 6.2
PPO + GRU
9.5 + 0.2
5.4
0.3
PPO
9.1 + 0.2
5.6
+ 0.1
Table 10: Ablation study of ExpGen on Maze. The table shows testing scores of networks trained on
200 maze instances. We present comparison between the proposed approach and LEEP, PPO+GRU and PPO, as well as an alternative ensemble policy with random actions upon ensemble disagreement. The mean and standard deviation computed using 10 runs with different seeds.
F.2 Evaluation of Various Similarity Measures for maxEnt
Tables and 12 present the results of our evaluation ExpGen equipped with a maxEnt exploration policy that uses either the Lo Le norms. The experiment targets the Maze and Heist environments and uses the same train and test procedures as in the main paper (25M training steps, score mean and standard deviation measured over 10 seeds).
Game
ExpGen Lo (Train)
ExpGen L2 (Train)
PPO (Train)
Heist
9.4 0.3
9.4 + 0.1
6.1 +
Maze
9.6
9.6
9.1 + 0.2
Table 11: Train scores of ExpGen using maxEnt policy with either Lo or L2 compared with PPO. The mean and standard deviation are measured over 10 seeds.
Game
ExpGen Lo (Test)
ExpGen La (Test)
PPO (Test)
Heist
7.4 +
2.4 + 0.5
Maze
8.2
8.3
5.6 0.1
Table 12:
Test scores of ExpGen using maxEnt policy with either Lo L2 compared with PPO. The mean and standard deviation are measured over 10 seeds.
The results demonstrate that both L2 and Lo allow ExpGen to surpass the PPO baseline for Maze and Heist environments, in which they perform similarly well at test time. This indicates that both are valid measures of state similarity for the maxEnt policy.
G Sample Complexity
One may wonder whether the leading approaches would benefit from training on additional environ- ment steps. A trained agent can still fail at test time either due to poor generalization performance (overfitting on small number training domains) or due insufficient training steps of the policy (underfitting). In this work, we are interested in the former and design our experiments such that no method underfits. Figure shows IDAAC training for 100M steps on Maze and Jumper, illustrating that the best test performance is obtained at around 25M steps, and training for longer does not contribute further (and can even degrade performance). Therefore, although ExpGen requires more environment steps (on account of training its ensemble of constituent reward policies), training for longer does not place our baseline (IDAAC) at any sort disadvantage.
22

---

Maze 100M steps (IDAAC)
Maze 100M steps (IDAAC)
10
10
test mean mean
test median median
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1cB
1e8
Jumper 100M steps (IDAAC)
Jumper 100M steps (IDAAC)
10
mean
test median
mean
train
0.75
1,00
0.00
0.25
1.00
166
168
Figure 13: The mean and median of the accumulated reward for IDAAC trained for 100M steps, averaged over 10 runs with different seeds. The curves show that test-reward stagnates and even decreases beyond 25M steps.
23


Reviews:
[{"review_id": "jd7LoOHUKe", "reviewer": "Reviewer_LrJL", "summary": "This work studies generalization on unseen similar tasks, in a zero-shot manner, and discusses how invariance based approach to overfitting might not work all the time. The algorithm proposed, called ExpGen, has one part that explores the space, while a ensemble of agents are trained to do the reward optimization, and it claims to achieve sota results on ProcGen.", "questions": "- Do all methods, including baselines, go through a proper hyperparameter tuning and search phase conducted by the authors?", "limitations": "The authors do discuss the limitations in their work.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well written, and it introduces the problem statement and the challenges faced by the current methods well.\n2. There is a good discussion about the related work as well, which helps place the proposed method in a proper context.\n3. The set of experiments and tables are clear and the authors have given proper pointers to architecture choices and hyperparameters.\n4. The limitations section is also addressed very well. The work talks about the game of Dodgeball where all the methods suffer and some possible insights on things that could be looked into for this.\n5. The idea looks promising and can help in adding a useful contribution to the community.", "weaknesses": "1. Adding training progress plots can show how the evaluation scores evolved. \n2. Figure 4 is very difficult to read. It can also benefit from the more useful captions.\n3. The proposed method does not do well in all the games. Other than the invariance, is there anything else that might be a reason for this? And is there a study on improving on these games?\n4. For some baselines, the scores from the respective papers seem to have been used. It is possible the methods missed a proper hyperparameter tuning. Clarifying this in the work  and doing a proper hyperparameter search and tuning for all methods equally will be helpful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Training progress:\nExpGen combines already trained exploration driven policy with an ensemble of (trained) reward policies at test-time. Therefore the evaluation targets test environments and does not produce figures of the training progress.\nWe will clarify this point, as well as clarify Figure 4 and its caption in the final version.\n\nExpGen performance: \nIn the paper, ExpGen shows a notable benefit in some environments but not in all. We suspected it is due to the limitations of PPO which forms the ensemble of ExpGen. This is validated in our experiment described in the general comment (please see Fig. 1) where we evaluate ExpGen with an IDAAC ensemble, leading to high performance in all environments.\n\nHyperparameters for baselines: this is a good point. For LEEP, we performed our own hyperparameter search, following the advice of the paper's first author (we directly corresponded with him). For some domains, however, we could not recover the performance reported in the LEEP paper, and for those domains, we used LEEP's reported numbers, giving an advantage to LEEP. For IDAAC, we used the published code, and corresponded with the first author to verify that we are using the best hyperparameters (she conducted extensive tuning). Again, for some domains we could not reproduce her results, and in those cases used IDAAC's reported scores, giving IDAAC an advantage."}, {"review_id": "ffNURGoiQE", "reviewer": "Reviewer_gmPQ", "summary": "This paper studies zero-shot generalization in RL. They first make an interesting observation: that intrinsic novelty-based rewards corresponding to maximum entropy exploration exhibit a smaller generalization gap than extrinsic environment rewards on ProcGen games. This suggests that MaxEnt rewards are in some sense richer and harder to memorize, and the agent is less prone to overfitting to them and more likely to generalize its exploratory behavior to new levels. The MaxEnt reward here is implemented using a kNN based method, similar to ProtoRL of Yarats et al. \n\nBased on this observation, the paper proposes a new algorithm (ExpGen), which trains a MaxEnt exploratory policy (which, as noted above, should generalize its exploratory behavior well), as well as an ensemble of K exploitation policies using the usual extrinsic reward. At test time, if the ensemble agrees on an action (indicating good generalization), then this action is executed. If not (indicating poor generalization), the exploratory policy (which is assumed to generalize well) is executed instead for a certain number of steps, and the process is repeated. \n\nThis algorithm is evaluated on the ProcGen benchmark, where it is compared to a number of other published methods (PPO, PLR, UCB-DrAC, PPG, IDAAC, LEEP). On two games (Maze and Heist), it significantly outperforms the other methods. However, it significantly underperforms in several others. \n\nOverall, this paper has several things going for it: the insight that MaxEnt reward has more favorable generalization properties is definitely interesting, and I think the ExpGen algorithm (or some variant of it) has potential. However, the experiments do not (yet) convincingly make a case for this algorithm: while it shows advantages in some games, it significantly underperforms in others and its aggregate performance does not seem favorable. This may be fixable — as the authors note, the invariances induced by other algorithms are orthogonal to the contributions of ExpGen, so I suspect the benefits could be combined. I would suggest combining the architectural modifications and auxiliary losses from IDAAC with ExpGen - if this indeed combines the best of both algorithms, the resulting method would be a lot more convincing. This is discussed in the paper, but not done, and in my opinion needs to be tested. \n\nSecond, there are a number of presentation and/or methodological issues which also need to be addressed. Please see my comments in the Weaknesses section. \n\nI think that if all these issues can be addressed, then this would make a strong submission. I think the substance of this paper is good, but it probably requires another revision cycle to be polished enough for publication. \n\n**Post rebuttal update: the authors have addressed my two main concerns. They have shown that ExpGen can be combined with IDAAC to get robust and SOTA results across all ProcGen games, and they have shown that the baselines do not improve when given a larger sample budget. Based on this, and assuming the authors will also improve the presentation as promised, I think this is now a strong submission and have raised my score to a 7 (Accept).**\n", "questions": "My suggestions for improving the paper are as follows:\n\nMajor:\n\n- Try to improve the performance of ExpGen by combining it with algorithmic elements from IDAAC or other methods. As mentioned in the paper, these improvements are orthogonal so hopefully it should be possible — however, it’s not sufficient to hypothesize that the two can be combined, this needs to be validated experimentally. \n- It may be that ProcGen isn’t the best experimental testbed to showcase the strengths of the proposed algorithm (indeed, several of the games have dense rewards, and exploration bonuses can sometimes actually hurt performance in such settings). An alternative could be MiniHack (https://arxiv.org/abs/2109.13202) — these environments are also procedurally generated, and additionally for the most part have sparse rewards. You can easily replicate the current setup where there is a limited number of seeds at training, and they are also very fast to run. \n- Another potential testbed is the Habitat embodied AI environment. That has a limited number of training levels and overfitting/generalization is a severe issue. In particular, Object Navigation is currently not solvable via RL and current SOTA requires large-scale IL using expensive human demonstrations (https://arxiv.org/abs/2204.03514) (in particular, because they perform exploration which the agent also needs to do at test time). I think this could be a good fit for the proposed method, and if you could get it to work there that would be very convincing. \n- Tables 1 and 2 are hard to read due to the many environments and methods. I would suggest displaying aggregate metrics using the RLiable library (https://github.com/google-research/rliable). This will also determine if differences are statistically significant. \n- I’m concerned the comparison to baselines might not be fair. It is reported that each method trains for 25M steps. However, ExpGen trains not 1 but (K+1) policies. Are each of these trained for 25M steps, or are they trained for 25M / (K+1) steps? To fairly compare to a policy trained for 25M steps, it should be the latter. Please clarify. If each policy was trained for 25M steps, then the baselines should be trained for 25M * (K+1) steps. \n\n\n\n\nMinor:\n\n- The legend in Figure 3 is quite small and hard to find, which can make the reader confused as to what they’re looking at. I would suggest making it bigger and placing it below the subplots, instead of inside the first one. \n- The notation in Section 4.1 is confusing: specifically, the Latex rendering of “k-NN” make it look like “k minus NN”. I would suggest using “x^\\mathrm{kNN}” or something else that doesn’t use the minus sign. \n- Figure 5: it is not clear from the caption which task this is for. Please add this to the caption. \n- The references are currently missing from the main paper (although they show up in the version in the supplement), please fix. ", "limitations": "Yes. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses an important problem - zero-shot generalization in RL is relatively understudied compared to the singleton MDP setting, but important for many realistic settings\n- As mentioned above, the insights are original, and the algorithm (also original) follows nicely from them", "weaknesses": "- As mentioned above, the experiments in their current form are not convincing enough\n- There is potentially an important methodological issue which needs clearing up, namely ExpGen appears to use many more samples than the other algorithms it is compared to due to the use of several policies. \n- The presentation could be improved", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Methodological issue - sample size:\nThank you for raising this point. We address it in the general comment (accompanied with Rebuttal Figure 2). \n\nCombining ExpGen+IDAAC:\nThis is a valuable insight, one that is expressed by the other reviewers as well. In the experiments described in the general comment (accompanied by Rebuttal Figure 1) we address this point and show state-of-the-art performance in all ProcGen games.\n\nMinihack/Habitat: \nThank you for this suggestion. We will assess how to extend our evaluation to these domains. That said, we believe our results on ProcGen, including the new results in the rebuttal, already establish that ExpGen makes important progress in zero-shot generalization.\n\nRegarding Tables 1 and 2, we'll clarify them and also include aggregate metrics in the revised version. "}, {"review_id": "tpzhLwxVxy", "reviewer": "Reviewer_W3Xj", "summary": "The paper is based upon the key insight, that maxEnt exploratory policies exhibit a much smaller generalization gap than usual reward seeking policies.\nPrevious work introduced the framework of epistemic POMDPs, where a random action is chosen until the uncertainty of the policy is low enough again (quantified by policy ensemble members that agree or disagree).\nThe paper improves upon the previous work by using the maxEnt exploratory policy instead of the random policy when policy ensemble members disagree.\nThis allows to substantially improve generalization performance on ProcGen tasks previous methods fail on (Maze and Heist).", "questions": "- How do you select the k-NN? Is it also done using the L0 norm?\n- What is your hypothesis why the L0 works better?\n- How is it justified to use the L0 norm instead of the L2 norm for the k-NN entropy estimator? I.e. the particle based derivation will not work out.\n- The main point of the paper is, that maxEnt exploration policies transfer better in zero-shot generalization. This should be true irrespective of the final performance.\nThe investigation of this property (Fig. 3 and 6) therefore seems limited, as those are the environments the algorithm actually improves most on. I would raise my score, if the same evaluation could be provided for all environments investigated in the main experiments to get a better view of the empirical significance of this finding.\n- How are scores of the main experiments in table 1 and 2 computed?", "limitations": "The authors openly discuss limitations of their approach for other environments of the ProcGen suite, which are better suited for different algorithms that incorporate invariances into the trained policy.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Very well written introduction and related work.\n- Well executed idea to improve weaknesses of previous state-of-the-art algorithm for specific generalization problems of the ProcGen suite.\n- Interesting and well executed empirical investigation.", "weaknesses": "**Major Weaknesses**\n\n- The usage of L0 norm instead of L2 norm in Eq. (8) is not justified, as to the best of our knowledge, the (log of the) k-NN distance does not approximate the entropy under this norm, but only under the euclidean (L2) norm. Therefore, maximising Eq. (8) might not yield a maxEnt policy.\n- Evaluation protocol of main experiments is not clear (see questions).\n\n**Minor Weaknesses**\n\n- In line 72 / 73, it is stated that the exploration policy is GUARANTEED to generalize. The paper does not contain a formal proof, but rather an empirical investigation, therefore this claim should be down-toned to reflect this correctly.\n\n**Remarks**\n\n- Please check the citation style again, often \\cite is used instead of \\citep.\n- Use \\eqref to reference equations.\n- Check notation (e.g. the true state as part of the history in line 128)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "L2/L0: We used L0 with an intuition that pixel changes are what matters for games like ProcGen, where the agents are localized. This is indeed an approximation, but one that we found to work well. We started an experiment with L2, but did not yet get the results. We remark that L2 is also an approximation (yet indeed a more justified one). We will clarify this point in the text.\n\nRegarding lines 72/73: we agree and will tone down this statement\n\nWe are thankful for your inclination to raise your score.\nWe are working on evaluating the generalization ability of MaxEnt exploration policy on all environments, and expect to finish the experiments during the discussion phase.\n\nRegarding the score in Tables 1 and 2: They present the average cumulative reward achieved on the train and test environments. Note that for the train environments, only 200 seeds (levels) are used, and for test environments, a new seed is sampled at every episode (from a pool of all possible seeds)."}, {"review_id": "BfzoxIbDH5", "reviewer": "Reviewer_VwcD", "summary": "---\nI have raised my score based on the answers and results provided by the authors during the rebuttal. \n\n---\nThis paper proposes an algorithm named ExpGen that can selectively exhibit a maximum entropy exploration behavior at test time by measuring epistemic uncertainty through ensemble of policies. In order to obtain a policy trained to maximize the entropy, the entropy is formulated as an intrinsic reward, where the sample based approximation to state distribution entropy is obtained using a trajectory of states are used as neighbors. The experimental results show that ExpGen is able to achieve highest score in two of the five ProcGen environments used in the paper, namely, Maze and Heist.", "questions": "- Why is $k$=1 a good choice (neighbor size)? From the formulation, it seems having just one sample would have a large bias, and thus more samples would lead to better performance.\n- The variable $k$ is used twice to represent the majority volume in ensemble. The values seem to largely vary from task to task. Was there a pattern or intuition behind selecting a good candidate for a given task?\n- Using difference in RGB images as the difference in states does not quite make sense. Since states within the same trajectory are treated as nearest neighbors, I do not see a clear connection between states that are close in time and close in image space. Shouldn't two have a correlation in order for the entropy estimation to make sense? For example, it would make more sense if the observations are mapped to an embedding space that clusters temporally near states and the embeddings are used to calculate the L2 norm between states.\n- Considering the computational burden of training ensemble of policies, calculating L2 distance between two images of size 64x64x3 does not seem like a much of a burden. Was there a large time gap when full original image size was used? How does the performance differ when the original image is used versus different down sampling strategies?\n- is $i$ randomly selected?\n- Regarding the experiment done in Figure 5, expanding the observed pattern, it seems that lower $\\gamma$ and lower $\\beta$ may enable PPO to achieve a score similar to ExpGen. Is there an ablation of such?\n- While the algorithm is designed upon the idea of preference of generalizing over different tasks (as described  by $\\mathcal{L}_{emp}$) compared to inductive biases, the experimental results does not seem to reflect this argument as the performance increase is seen only in a small subset of evaluation tasks. \n- As mentioned in 6.1, is there a result of using IDAAC as the base model for ExpGen?\n\nMinor remarks.\n- In Equation 5, it is not specified how and where $x_i$ is sampled from.\n- Figure 6, legend labels GPU, which should be GRU\n- Reference section is missing in the paper (but included in supplementary)", "limitations": "Although limitations are included as part of the discussion, some clear limitation could be directly stated, which should help making a clear distinction of ExpGen to other algorithms. For example, comparing memory and computational efficiency against using designs like PPO with extrinsic and intrinsic reward instead of ExpGen with ensemble.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper proposes novel framework that combines the idea of switching strategies based on the epistemic uncertainty and leveraging the maximum entropy policy to bridge the gap between the performances during train and test time.", "weaknesses": "Some parts of the logic are unclear and there are few typos and mistakes in the paper. The experimental draws some what countering argument against the motivation of using the general framework instead of ones based on inductive biases, e.g., IDAAC. This naturally leads to combination of IDAAC and ExpGen, but is mentioned in the paper, but not experimented. Another weakness would be the memory and computational inefficiency, resulting from needing to train and inferencing ensemble of policies. Finally, the empirical studies does not draw a clear picture of why and how ExpGen works, i.e., most results are observed performances of the comprehensive algorithm, lacking a detailed ablation of each components of the algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We include an evaluation of ExpGen+IDAAC (described in the general comment and accompanied by Rebuttal Figure 1), showcasing state-of-the-art performance on all ProcGen environments.\n\nRegarding memory and computational inefficiencies: \nThis is a valid point, which we address in the general comment (accompanied by Rebuttal Figure 2).\n\nAblation of ExpGen components:\nThank you for this feedback, we will add ablation experiments to the final version. We also wish to refer you to the appendix (section D) for the ablation study.\n\nRegarding the choice of the value of k in KNN, we conducted hyperparameter tuning for various values - based on which k=1 yielded the best performance. We will include the evaluation in the final version of this paper. \n\nThe majority volume of the ensemble k:\nWe treat k as a hyperparameter that we tune. Please see Table 10 in the appendix (section D) for an example of tuning k for the maze environment based on an evaluation for various values of k.\n \nIn this work we implemented KNN from the RGB image as a way to discern novel states. This was achieved without learned embeddings, but we agree with you that encoding temporally-near states in embedding space has the potential to further improve entropy-maximization performance.\n\nWith regards to down-sampling the observed state (full image), our implementation of KNN uses down-sampling primarily in order to smooth out noisy pixels in the observation, as well as extract a more concise state representation that is less sensitive to minor changes. It also helped reduce the training time of the entropy-maximizing policy (MaxEnt) by 50\\%.\n\nRegarding an evaluation of lower values of alpha and beta:\nFigure 5 illustrates that even the best values are still significantly worse compared to ExpGen. That said, we will include an evaluation of lower values of alpha and beta in the final version.\n\nThank you for your insight on combining ExpGen+IDAAC. We include this experiment in the rebuttal: described in the general comment and accompanied with Rebuttal Figure 1. The ExpGen+IDAAC algorithm achieves state-of-the-art results in high-difficulty environments and yields on-par performance in the remaining environments."}]

Author rebuttal:
Thank you for your valuable insights and suggestions. 
This paper is the first to incorporate exploration driven behavior at test-time towards generalization in RL, and in doing so, achieves state-of-the-art performance in environments that are widely regarded as challenging by all the leading algorithms (e.g. Maze, Heist, Jumper). We are confident that the research community would benefit from this work, due to the significance of zero-shot generalization for RL, and would attract further research into improving our understanding of the role of exploration towards generalization.

We first address all reviewers and describe new results based on the reviewers' suggestions.

In our submission, ExpGen forms an ensemble of reward driven PPO agents. On its own, this approach surpasses the state-of-the-art
in several challenging ProcGen environments (e.g. Maze, Heist, Jumper).
Per the reviewers' recommendation, we combine ExpGen+IDAAC to evaluate variant in which IDAAC reward policies comprise the ensemble (rather than PPO) since IDAAC is the SOTA in all remaining ProcGen environments (e.g., Plunder, Miner, BigFish, etc.) - please see rebuttal Fig. 1.
ExpGen+IDAAC outperforms IDAAC alone in several games (still very significant on Heist, Maze, and Jumper) and performs on-par across all others. **We believe these are strong results**. Moreover, this demonstrates that the applicability of ExpGen is not limited to PPO, and that ExpGen can be applied in the future to other, more powerful reward seeking models to infuse them with exploration driven behavior at test time. 

Sample complexity: Several reviewers raised the concern that ExpGen benefits from training on more environment steps (because of exploration + ensemble policies). Let us explain why this is not the case: An agent can fail at test time either due to bad generalization performance (overfitting, e.g., due to a small number of training domains), or due to insufficient training steps of the policy (underfitting). In this work we are interested in the former, and design our experiments such that no method underfits. The Rebuttal Figure 2 shows IDAAC training for 100M steps, which demonstrates that the best test performance is obtained at around 25M steps, and training for longer does not help (and can even degrade performance). Thus, while it is true that our method requires more samples, **our baselines are not in any disadvantage**. Adding constraints on sample complexity to ExpGen is interesting, but is out of scope for this study, which focuses on unlimited samples, but a very limited set of 200 training levels.

Meta-review: This paper proposes a new method for zero-shot generalisation in RL, where the main idea is to train an ensemble of policies using extrinsic rewards along with an exploration policy using an intrinsic reward for maximum state entropy, and use either of them to act at test time, depending on how much the ensemble of policies agree with each other. The majority of the reviewers agreed that the proposed method is novel and sensible. Some of the reviewers had concerns around the poor performances on some tasks and the justification of the choice of distance metric for the kNN method. However, the authors provided much stronger results and conducted additional ablations for kNN during the rebuttal period, and the reviewers were satisfied with the new results. One of the reviewers who expressed a similar concern did not follow up during the discussion period. Since the rest of the reviewers said the new results addressed their concerns, I recommend accepting this paper.
Decision: Accept (poster)