SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ev Zisselman*, Itai Lavie, Daniel Soudry, Aviv Tamar
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 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 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.
Recent developments in reinforcement learning (RL) led to algorithms that surpass human experts in a broad range of tasks [Mnih et al., 2015, Vinyals et al., 2019, Schrittwieser et al., 2020, Wurman et al., 2022]. In most cases, the RL agent is tested on the same task it was trained on, and is not guaranteed to perform well on unseen tasks. In zero-shot generalization for RL (ZSG-RL), however, the goal is to train an agent on training domains to act optimally in a new, previously unseen test environment [Kirk et 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. 1a, or background image in Jumper, Fig. 1c) and dynamics (e.g., wall positions in Maze, Fig. 1d, and key positions in Heist, Fig. 1e).
Previous studies focused on identifying various invariance properties in the tasks, and designing corresponding invariant policies, through an assortment of regularization and augmentation techniques [Igl et al., 2019, Cobbe et al., 2019, Wang et al., 2020, Lee et al., 2019a, Raileanu et al., 2021, Raileanu and Fergus, 2021, Cobbe et al., 2021, Sonar et al., 2021, Bertran et al., 2020, Li et al., 2021]. For example, a policy that is invariant to the color of agents is likely to generalize well in BigFish. More intricate invariances include the order of observations in a trajectory [Raileanu and Fergus, 2021], and the length of a trajectory, as reflected in the value function [Raileanu and Fergus, 2021].
*Correspondence E-mail: ev_zis@campus.technion.ac.il
(a) BigFish
(b) Plunder
(c) Jumper
(d) Maze
(e) Heist
(f) Hidden Maze
Figure 1: (a), (b), (c), (d) 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 experiment<sup>2</sup>. Imagine Maze, but with the walls and goal hidden in the observation (Fig. 1f). Arguably, this is the most task-invariant observation possible, such that a solution can still be obtained in a 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 in which training maze it is currently in. Obviously, such a 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-shot 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 an unobserved variable. The resulting decision making problem, termed the epistemic POMDP, may require some exploration at test time to resolve uncertainty. Ghosh et al. [2021] 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 we train an agent to explore the training domains using a maximum entropy objective [Hazan et al., 2019, Mutti et al., 2021], we observe that the learned exploration behavior generalizes surprisingly well—much better than the generalization attained when training the agent to maximize reward. Intuitively, this can be explained by the fact that reward is a strong signal that leads to a specific behavior that the agent can 'memorize' during training, while exploration is naturally more varied, making it harder to memorize and overfit.
Exploration by itself, however, is not useful for solving new tasks. Our algorithm, Explore to Generalize (ExpGen), additionally trains an ensemble of reward-seeking agents. At test time, either the ensemble agrees on an action, and we generalize well, or we take exploratory actions using the exploration policy, which we demonstrate to generalize, and drive us to a novel part of the state space, where the ensemble may potent
Normalized test performance on five ProcGen games
Figure 2: Normalized test Performance for ExpGen, LEEP, IDAAC, DAAC, and PPO, on five ProcGen games. ExpGen shows state-of-the-art performance on test levels of Maze, Heist and Jumper; games that are notoriously challenging for other leading approaches. The scores are normalized as proposed by [Cobbe et al., 2020].
ExpGen is simple to implement, and can be used with any reward maximizing RL algorithm. Combined with vanilla PPO, ExpGen significantly improves the state-of-the-art (SOTA) on several ProcGen games for which previous methods fail (see Fig. 2). ExpGen also significantly improves
2We validated this experiment empirically, using a recurrent policy on Maze with 128 training tasks, observing  $85\%$  success rate on training domains, and test success similar to a 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 state-of-the-art achieved  $66\%$ . When combined with IDAAC [Raileanu and Fergus, 2021], the leading invariance-based algorithm, ExpGen achieves state-of-the-art performance on the full ProcGen suite (the full results are provided in Appendix D).
Generalization in RL The recent survey by Kirk et al. [2021] provides an extensive review of generalization in 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., 2020, Lee et al., 2019a, Raileanu et al., 2021], task interpolation [Yao et al., 2021] or various regularization technique, such as dropout [Igl et al., 2020] and batch normalization [Farebrother et al., 2018, Igl et al., 2020]. Leading approaches, namely IDAAC [Raileanu and Fergus, 2021] and PPG [Cobbe et al., 2021], investigate the advantages of decoupling policy and value functions for generalization, whereas Jiang et al. [2021] propose automatic curriculum learning of levels.
A different approach is to add inductive bias to the neural network policy or learning algorithm. Approaches such as Tamar et al. [2016], Vlastelica et al. [2021], 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., 2020] combine learning with classical graph planning to generalize across various planning domains. These approaches require some knowledge about the problem structure (e.g., a 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 et al., 2021, Bertran et al., 2020, Li et al., 2021, Igl 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 an ensemble of policies, each on a separate subset of the training environment, with a loss function that encourages agreement between the ensemble members. Effectively, the KL 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 a policy that is explicitly trained to explore, such as a maximum-entropy policy. Consequently, we observe that our approach leads to significantly better performance at test time.
State Space Maximum Entropy Exploration Maximum entropy exploration (maxEnt, Hazan et al. 2019, Mutti et al. 2021) is an unsupervised learning framework that trains policies that maximize the entropy of their state-visitation frequency, leading to a behavior that continuously explores the environment state space. Recently, maximum entropy policies have gained attention in RL [Liu and Abbeel, 2021b,a, Yarats et al., 2021, Seo et al., 2021, 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 a 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 to the same environments during pre-training and test phases, with the only distinction being the lack of extrinsic reward during pre-training. 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 on maxEnt exploration.
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, A, O, P_{init}, P, \Sigma, r, \gamma)$ , where  $S \in \mathbb{R}^{|S|}$  and  $A \in \mathbb{R}^{|A|}$  are the state and actions spaces,  $O$  is the observation space,  $P_{init}$  is an initial state distribution,  $P$  is the transition kernel,  $\Sigma$  is the observation function,  $r: S \times A \to \mathbb{R}$  is the reward function, and  $\gamma \in [0,1)$  is the discount factor. The agent starts from initial state  $s_0 \sim P_{init}$  and at time  $t$  performs an action  $a_t$  on the environment that yields a reward  $r_t = r(s_t, a_t)$ , and an observation  $o_t = \Sigma(s_t, a_t) \in O$ . Consequently, the environment transitions into the next state according to  $s_{t+1} \sim P(\cdot | s_t, a_t)$ . Let the history at time  $t$  be  $h_t = \{o_0, a_0, r_0, o_1, a_1, r_1, \dots, o_t\}$ , the sequence of observations, actions and rewards. The agent's next action is outlined by a policy  $\pi$ , which is a stochastic mapping from the history to an action probability  $\pi(a|h_t) = P(a_t = a|h_t)$ . In our formulation, a history-dependent policy (and not a 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].
Zero-Shot Generalization for RL
We assume a prior distribution over POMDPs  $P(M)$ , defined over some space of POMDPs. For a given POMDP, an optimal policy maximizes the expected discounted return  $\mathbb{E}_{\pi,M}[\sum_{t=0}^{\infty} \gamma^{t} r(s_{t},a_{t})]$ , where the expectation is taken over the policy  $\pi(h_{t})$ , and the state transition probability  $s_{t} \sim P$  of POMDP  $M$ . Our generalization objective in this work is to maximize the discounted cumulative reward taken in expectation over the POMDP prior, also termed the population risk:
$$
\mathcal {R} _ {p o p} (\pi) = \mathbb {E} _ {M \sim P (M)} \left[ \mathbb {E} _ {\pi , M} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r \left(s _ {t}, a _ {t}\right) \right] \right]. \tag {1}
$$
Seeking a policy that performs well in expectation over any POMDP from the prior corresponds to zero-shot generalization.
We assume access to  $N$  training POMDPs  $M_{1},\ldots ,M_{N}$  sampled from the prior,  $M_{i}\sim P(M)$ . Our goal is to use  $M_1,\dots ,M_N$  to learn a policy that performs well on objective 1. A common approach is to optimize the empirical risk objective:
$$
\mathcal {R} _ {e m p} (\pi) = \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbb {E} _ {\pi , M _ {i}} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r (s _ {t}, a _ {t}) \right] = \mathbb {E} _ {M \sim \hat {P} (M)} \left[ \mathbb {E} _ {\pi , M} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r (s _ {t}, a _ {t}) \right] \right], \tag {2}
$$
where the empirical POMDP distribution can be different from the true distribution, i.e.  $\hat{P}(M) \neq P(M)$ . In general, a policy that optimizes the empirical risk (Eq. 2) may perform poorly on the population risk (Eq. 1)—this is known as overfitting in statistical learning theory [Shalev-Shwartz and Ben-David, 2014], and has been analyzed recently also for RL [Tamar et al., 2022].
Maximum Entropy Exploration
In the following we provide the definitions for the state distribution and the maximum entropy exploration objective. For simplicity, we discuss MDPs—the fully observed special case of POMDPs where  $O = S$ , and  $\Sigma(s, a) = s$ .
A policy  $\pi$ , through its interaction with an MDP, induces a  $t$ -step state distribution  $d_{t,\pi}(s) = p(s_t = s|\pi)$  over the state space  $S$ . Let  $d_{t,\pi}(s,a) = p(s_t = s,a_t = a|\pi)$  be its  $t$ -step state-action counterpart. For the infinite horizon setting, the stationary state distribution is defined as  $d_{\pi}(s) = \lim_{t\to \infty}d_{t,\pi}(s)$ , and its  $\gamma$ -discounted version as  $d_{\gamma ,\pi}(s) = (1 - \gamma)\sum_{t = 0}^{\infty}\gamma^{t}d_{t,\pi}(s)$ . We denote the state marginal distribution as  $d_{T,\pi}(s) = \frac{1}{T}\sum_{t = 0}^{T}d_{t,\pi}(s)$ , which is a marginalization of the  $t$ -step state distribution over a finite time  $T$ . The objective of maximum entropy exploration is given by:
$$
\mathcal {H} (d (\cdot)) = - \mathbb {E} _ {s \sim d} [ \log (d (s)) ], \tag {3}
$$
where  $d$  can be regarded as either the stationary state distribution  $d_{\pi}$  [Mutti and Restelli, 2020], the discounted state distribution  $d_{\gamma,\pi}$  [Hazan et al., 2019] or the marginal state distribution  $d_{T,\pi}$  [Lee et al., 2019b, Mutti and Restelli, 2020]. In our work we focus on the finite horizon setting and
adapt the marginal state distribution  $d_{T,\pi}$  in which  $T$  equals the episode horizon  $H$ , i.e. we seek to maximize the objective:
$$
\mathcal {R} _ {\mathcal {H}} (\pi) = \mathbb {E} _ {M \sim \hat {P} (M)} [ \mathcal {H} (d _ {H, \pi}) ] = \mathbb {E} _ {M \sim \hat {P} (M)} \left[ \mathcal {H} \left(\frac {1}{H} \sum_ {t = 0} ^ {H} d _ {t, \pi} (s)\right) \right], \tag {4}
$$
which yields a 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 et al., 2019, Lee et al., 2019b]. More recently, a branch of algorithms that employ non-parametric entropy estimation [Liu and Abbeel, 2021a, Mutti et al., 2021, Seo et al., 2021] has emerged, circumventing 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 al., 2003], as elaborated in the next section.
4 The Generalization Ability of Maximum Entropy Exploration
In this section we present an empirical observation—policies trained for maximum entropy exploration (maxEnt policy) generalize well. First, we 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 et al., 1997, Singh et al., 2003], as described here. Let  $X$  be a random variable over the support  $X \subset \mathbb{R}^m$  with a probability mass function  $p$ . Given the probability of this random variable, its entropy is obtained by  $\mathcal{H}_X(p) = -\mathbb{E}_{x \sim p}[\log(p)]$ . Without access to its distribution  $p$ , the entropy can be estimated using  $N$  samples  $\{x_i\}_{i=1}^N$  by the  $k$ -NN estimator Singh et al. [2003]:
$$
\hat {\mathcal {H}} _ {X} ^ {k, N} (p) \approx \frac {1}{N} \sum_ {i = 1} ^ {N} \log \left(\left\| x _ {i} - x _ {i} ^ {k - \mathrm {N N}} \right\| _ {2}\right), \tag {5}
$$
where  $x_{i}^{k\text{-NN}}$  is the  $k$ -NN sample of  $x_{i}$  from the set  $\{x_{i}\}_{i = 1}^{N}$ .
To estimate the distribution  $d_{H,\pi}$  over the states  $S$ , we consider each trajectory as  $H$  samples of states  $\{s_t\}_{t=1}^H$  and take  $s_t^{k-NN}$  to be the  $k$ -NN of the state  $s_t$  within the trajectory, as proposed by previous works (APT, Liu and Abbeel [2021a], RE3, Seo et al. [2021], and APS, Liu and Abbeel [2021a]),
$$
\hat {\mathcal {H}} ^ {k, H} \left(d _ {H, \pi}\right) \approx \frac {1}{H} \sum_ {t = 1} ^ {H} \log \left(\left\| s _ {t} - s _ {t} ^ {k - \mathrm {N N}} \right\| _ {2}\right). \tag {6}
$$
Next, similar to previous works, since this sampled estimation of the entropy (Eq. 6) is a sum of functions that operate on each state separately, it can be considered as an expected reward objective  $\hat{\mathcal{H}}^{k,H}(d_{H,\pi})\approx \frac{1}{H}\sum_{t = 1}^{H}r_I(s_t)$  with the intrinsic reward function:
$$
r _ {I} \left(s _ {t}\right) := \log \left(\left\| s _ {t} - s _ {t} ^ {k - \mathrm {N N}} \right\| _ {2}\right). \tag {7}
$$
This formulation enables us to 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  $t$  the state  $s_t^{k\text{-NN}}$  is chosen from previous states  $\{s_i\}_{i=1}^{t-1}$  of the same episode.
Another challenge stems from the computational complexity of calculating the  $L_{2}$  norm of the  $k$ -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  $s_{i}$  (i.e.  $64 \times 64$  RGB image), we sub-sample (denoted  $\downarrow$ ) the observation by applying average pooling of  $3 \times 3$  to produce an image  $s_{i}^{\downarrow}$  of size  $21 \times 21$ , resulting in:
$$
r _ {I} \left(s _ {t}\right) := \log \left(\left\| s _ {t} ^ {\downarrow} - s _ {t} ^ {k - \mathrm {N N}, \downarrow} \right\| _ {2}\right). \tag {8}
$$
(a) Entropy Reward
(b) Extrinsic Reward
We emphasize that we do not modify the termination condition of each game. However, a maxEnt policy will learn to avoid termination, as this increases the sum of intrinsic rewards. In Figure 3 we display the states visited by a maxEnt policy on Maze. We also experimented with  $L_{0}$  as the state similarity measure instead of  $L_{2}$ , which resulted in similar performance (see Appendix F.2).
4.2 Generalization of maxEnt Policy
The generalization gap describes the difference between the reward accumulated during training  $\mathcal{R}_{emp}(\pi)$  and testing  $\mathcal{R}_{pop}(\pi)$  of a 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 4: Generalization ability of maximum entropy vs. extrinsic reward: (a) Score of maximum entropy. (b) Score of extrinsic reward. Training for maximum entropy exhibits a small generalization gap in Maze, Jumper and Miner. Average and standard deviation are obtained using 4 seeds.
Figure 3: Example of a maxEnt trajectory on Maze. The policy visits 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 an '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. 8). 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] $^3$ .
The train and test return scores are shown in Fig. 4a. In all three environments, we demonstrate a small generalization gap, as test performance on 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 maxEnt policy exhibits a small generalization gap of  $8.5\%$  and  $4.3\%$ , respectively. In addition, we verify that the train results are near optimal by comparing with a 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<sup>4</sup>. The results for this experiment, shown in Fig. 4b, illustrate that the generalization gap for extrinsic reward
An extensive discussion on the importance of memory for the maxEnt objective is in Appendix B.3.
4We train for extrinsic reward using an architecture identical to that of the intrinsic reward, with the exception of the memory unit. Incorporating a 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 a large generalization gap for Maze (38.8%) and Jumper (27.5%), while Miner exhibits a moderate generalization gap of 13.1%. For an evaluation on all ProcGen games, please see Appendix B.1.
5 Explore to Generalize (ExpGen)