SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Matthew T. Jackson*  
University of Oxford
Minqi Jiang  
UCL
Jack Parker-Holder Google DeepMind
Risto Vuorio  
University of Oxford
Chris Lu  
University of Oxford
Gregory Farquhar  
Google DeepMind
Shimon Whiteson  
University of Oxford
Jakob N. Foerster  
University of Oxford
The past decade has seen vast progress in deep reinforcement learning (RL) on the back of algorithms manually designed by human researchers. Recently, it has been shown that it is possible to meta-learn update rules, with the hope of discovering algorithms that can perform well on a wide range of RL tasks. Despite impressive initial results from algorithms such as Learned Policy Gradient (LPG), there remains a generalization gap when these algorithms are applied to unseen environments. In this work, we examine how characteristics of the meta-training distribution impact the generalization performance of these algorithms. Motivated by this analysis and building on ideas from Unsupervised Environment Design (UED), we propose a novel approach for automatically generating curricula to maximize the regret of a meta-learned optimizer, in addition to a novel approximation of regret, which we name algorithmic regret (AR). The result is our method, General RL Optimizers Obtained Via Environment Design (GROOVE). In a series of experiments, we show that GROOVE achieves superior generalization to LPG, and evaluate AR against baseline metrics from UED, identifying it as a critical component of environment design in this setting. We believe this approach is a step towards the discovery of truly general RL algorithms, capable of solving a wide range of real-world environments.
Figure 1: Out-of-distribution performance on Atari--after meta-training exclusively on Grid-World levels, our method (GROOVE) significantly outperforms LPG on Atari. Improvement is measured as a percentage of mean human-normalized return over 5 seeds.
*Correspondence to jackson@robots.ox.ac.uk.
The past decade has seen vast progress in deep reinforcement learning [Sutton and Barto, 1998, RL], a paradigm whereby agents interact with an environment to maximize a scalar reward. In particular, deep RL agents have learned to master complex games [Silver et al., 2016, 2017, Berner et al., 2019], control physical robots [OpenAI et al., 2019, Andrychowicz et al., 2020, Miki et al., 2022] and increasingly solve real-world tasks [Degrave et al., 2022]. However, these successes have been driven by the development of manually-designed algorithms, which have been refined over many years to tackle new challenges in RL. As a result, these methods do not always exhibit the same performance when transferred to new tasks [Henderson et al., 2018, Andrychowicz et al., 2021] and are limited by our intuitions for RL.
Recently, meta-learning has emerged as a promising approach for discovering general RL algorithms in a data-driven manner [Beck et al., 2023b]. In particular, Oh et al. [2020] introduced Learned Policy Gradient (LPG), showing it is possible to meta-learn an update rule on toy environments and transfer it zero-shot to train policies on challenging, unseen domains. Despite impressive initial results, there remains a significant generalization gap when these algorithms are applied to unseen environments. In this work, we seek to learn general and robust RL algorithms, by examining how characteristics of the meta-training distribution impact the generalization of these algorithms.
Motivated by this analysis, our goal is to automatically learn a meta-training distribution. We build on ideas from Unsupervised Environment Design [Dennis et al., 2020, UED], a paradigm where a student agent trains on an adaptive distribution of environments proposed by a teacher, which seeks to propose tasks which maximize the student's regret. UED has typically been applied to train single RL agents, where it has been shown to produce robust policies capable of zero-shot transfer to challenging human-designed tasks. Instead, we apply UED to the meta-RL setting of meta-learning a policy optimizer, which we refer to as policy meta-optimization (PMO). For this, we propose algorithmic regret (AR), a novel metric for selecting meta-training tasks, in addition to a method building on LPG and ideas from UED. We name our method General RL Optimizers Obtained Via Environment Design, or GROOVE.
We train GROOVE on an unstructured distribution of Grid-World environments, and rigorously examine its performance on a variety of unseen tasks—ranging from challenging Grid-Worlds to Atari games. When evaluated against LPG, GROOVE achieves significantly improved generalization performance on all of these domains. Furthermore, we compare AR against prior environment design metrics proposed in UED literature, identifying it as a critical component for environment design in this setting. We believe this approach is a step towards the discovery of truly-general RL algorithms, capable of solving a wide range of real-world environments.
We implement GROOVE and LPG in JAX [Bradbury et al., 2018], resulting in a meta-training time of 3 hours on a single V100 GPU. As well as being the first complete and open-source implementation of LPG, we achieve a major speedup against the reference implementation, which required 24 hours on a 16-core TPU-v2. This will enable academic labs to perform follow-up research in this field, where compute constraints have long been a limiting factor.
Figure 2: GROOVE meta-training (left) and meta-testing (right). During meta-training, levels are sampled from both the level curator and sampler. Agents are trained by an optimizer for multiple updates, before a meta-optimizer updates the optimizer based on agent return. At the end of an agent's lifetime, its regret is calculated and the level curator is updated. During meta-testing, the trained optimizer is applied to previously unseen environments and agent architectures.
- In order to distinguish this problem setting from traditional meta-RL, we provide a novel formulation of PMO using the Meta-UPOMDP (Section 2.1).
- We propose AR (Section 3.2), a novel regret approximation for PMO, and GROOVE (Section 3.3), a PMO method using AR for environment design.
- We analyze how features of the meta-training distribution impact generalization in PMO (Section 4.2) and demonstrate AR as a proxy for task informativeness (Section 4.3).
- We extensively evaluate GROOVE against LPG, demonstrating improved in-distribution robustness and out-of-distribution generalization (Section 4.4).
- We perform an ablation of AR, demonstrating the insufficiency of existing methods (PLR and LPG) without AR, as well as the impact of the antagonist agent in AR (Section 4.5).
2 Problem Setting and Preliminaries
2.1 Formulating Policy Meta-Optimization for Environment Design
Policy Meta-Optimization In this work, we consider a subproblem of meta-RL which we refer to as policy meta-optimization (PMO). PMO is a bilevel optimization problem. In the inner loop, a collection of agents each interact with their associated environments and are updated with a policy optimizer. Following a series of inner-loop updates, the outer loop updates the policy optimizer in order to maximize the performance of these agents. PMO only trains the policy optimizer in the outer loop, while the initial agent parameters for each task are generated by a static initialization function.
Unsupervised Environment Design In UED [Dennis et al., 2020], a teacher is given the problem of designing an environment distribution which is maximally useful for training a student agent. UED formalizes this problem setting with the Underspecified Partially Observable Markov Decision Process (UPOMDP), an extension of the POMDP with additional free parameters  $\phi$  that parameterize those aspects of the environment which the teacher can modify throughout training. In prior work, this paradigm has been used to adapt environment distributions to facilitate the learning of a robustly transferable policy [Jiang et al., 2021b, Parker-Holder et al., 2022a]. However, in our problem setting of PMO, the central focus is on learning the update rule itself, with the goal of transferring the update rule to new environments.
Problem Formulation We therefore introduce the Meta-UPOMDP, an extension of the UPOMDP, to account for this difference. Formally, the Meta-UPOMDP is defined by the tuple  $\langle \mathcal{A},\mathcal{O},\mathcal{S},\mathcal{T},\mathcal{I},\mathcal{R},\gamma ,\Phi ,\Theta \rangle$  . The first components correspond to the standard UPOMDP, where  $\mathcal{A}$  is the action space,  $s$  is the state space,  $\mathcal{O}$  is the observation space, and  $\mathcal{T}:S\times \mathcal{A}\times \Phi \mapsto S$  is the transition function. Upon each transition, the student agent receives an observation according to the observation function  $\mathcal{I}:S\mapsto \mathcal{O}$  and a reward according to the reward function  $\mathcal{R}:S\times \mathcal{A}\mapsto \mathbb{R}$  . Here, the free parameters  $\phi \in \Phi$  control the variable aspects of the environment, such as the  $x,y$  -positions of obstacles in a 2D maze.1
The Meta-UPOMDP models PMO, in which an optimizer  $\mathcal{F}:\Theta \times \mathbb{T}\mapsto \Theta$  learns to update an agent's parameters  $\Theta$  given the sequence of states, actions, rewards and termination flags corresponding to the agent's past experience  $\mathbb{T}$  in the environment. This update is performed after every transition over the agent's lifetime of  $N$  environment interactions, resulting in a sequence of parameters  $(\theta_0,\dots ,\theta_N)$ . The Meta-UPOMDP extends the UPOMDP include agent parameters  $\Theta$  and appends the agent lifetime  $N$  to the free parameters  $\phi$ , making it a controllable feature of tasks.
For an initialization of agent parameters  $\theta_0$  and free parameters  $\phi$ , we define the value of the optimizer to be  $V_{\phi, \theta_0}(\mathcal{F}) = \mathbb{E}_{\pi_{\theta_N}}[\sum_t^\infty \gamma^t r_t]$ , which is the expected return of the trained agent  $\pi_{\theta_N}$  on the environment specified by  $\phi$  at the end of its lifetime. Given an optimizer  $\mathcal{F}_{\eta}$  with meta-parameters  $\eta$ , we reformulate the PMO objective from Oh et al. [2020] to
$$
\mathcal {L} (\eta) = \mathbb {E} _ {\phi \sim p (\phi)} \mathbb {E} _ {\theta_ {0} \sim p \left(\theta_ {0}\right)} \left[ V _ {\phi , \theta_ {0}} \left(\mathcal {F} _ {\eta}\right) \right], \tag {1}
$$
where  $p(\phi)$  and  $p(\theta_0)$  are distributions of free parameters and initial agent parameters.
<sup>1</sup>The terms "tasks" and "levels" are used interchangeably to refer to settings of free parameters.
2.2 Learned Policy Gradient
Learned Policy Gradient [Oh et al., 2020, LPG] is a PMO method which trains a generalization of the actor-critic architecture [Barto et al., 1983]. This replaces the critic with a generalization of value functions from RL, that we refer to as bootstrap functions. Whilst value functions are trained to predict the expected discounted return from a given state, bootstrap functions predict an  $n$ -dimensional, categorical bootstrap vector, the properties of which are meta-learned by LPG.
LPG uses a reverse-LSTM [Hochreiter and Schmidhuber, 1997] to learn a policy update for each agent transition, conditioned on all future episode transitions. For a single update to agent parameters  $\theta$  at time-step  $t$ , LPG outputs targets  $\hat{y}_t, \hat{\pi}_t = U_\eta(x_t | x_{t+1}, \ldots, x_T)$ , where  $x_t = [r_t, d_t, \gamma, \pi_\theta(a_t | s_t), y_\theta(s_t), y_\theta(s_{t+1})]$  is a vector containing reward  $r_t$ , episode-termination flag  $d_t$ , discount factor  $\gamma$ , probability of the chosen action  $\pi_\theta(a_t | s_t)$ , and bootstrap vectors for the current and next states  $y_\theta(s_t)$  and  $y_\theta(s_{t+1})$ . The targets  $\hat{y}$  and  $\hat{\pi}$  update the bootstrap function and policy respectively, giving the update rule
$$
\Delta \theta \propto \left[ \nabla_ {\theta} \log \pi_ {\theta} (a | s) \hat {\pi} - \alpha_ {y} \nabla_ {\theta} D _ {\mathrm {K L}} (y _ {\theta} \| \hat {y}) \right]. \tag {2}
$$
2.3 Learning Robust Policies via Minimax-Regret UED
A trivial example of UED is domain randomization [Jakobi, 1997, DR], in which the teacher generates an environment distribution by uniformly sampling free parameters from the UPOMDP. By sampling randomly, DR often fails to generate environments with interesting structure: they may be trivial, impossible to solve, or irrelevant to downstream tasks of interest. An alternative approach is to train an adversarial minimax teacher, whose objective is to generate environments which minimize the agent's return. This has the benefit of adapting the environment distribution to the agent's current capability, by presenting it with environments on which it performs poorly. However, by naively minimizing return, the teacher is incentivized to generate environments which are impossible for the agent to solve.
In contrast to this, Dennis et al. [2020] propose the minimax-regret objective for UED, in which the teacher aims to maximize the agent's regret, the difference between the achieved and maximum return. Unlike return minimization, this disincentivizes the teacher from generating unsolvable levels where regret would be 0. However, since computing the maximum return is generally intractable, methods implementing this objective have proposed approximations of regret. PAIRED [Dennis et al., 2020] co-trains an antagonist agent with the original (protagonist) agent, estimating regret as the difference in their performance. PLR [Jiang et al., 2021b,a] curates a buffer of high-regret levels, rather than training a generative model to produce them. In this, a range of regret approximations are evaluated, with positive value loss, L1 value loss, and maximum Monte Carlo achieving consistent performance across the evaluated domains.
3 Adversarial Environment Design for Policy Meta-Optimization
In this section, we introduce GROOVE, a novel method for PMO. We begin by formulating the minimax-regret objective implemented by GROOVE, followed by AR, a novel regret approximation designed for PMO. Finally, we discuss existing approaches to environment design and motivate the selection of a curation-based approach.
3.1 Minimax-Regret as a Meta-Objective
Our method trains an optimizer  $\mathcal{F}_{\eta}$  over an adversarial distribution of environments, in which the UED adversary's objective is to maximize the regret of the meta-learner, defined as
$$
\operatorname {R E G R E T} (\eta , \phi) = V _ {\phi} \left(\eta_ {\phi} ^ {*}\right) - V _ {\phi} (\eta). \tag {3}
$$
Here,  $\eta_{\phi}^{*}$  denotes the optimal meta-parameters for updating the student's policy on an environment instance  $\phi$ , i.e.,  $\eta_{\phi}^{*} = \arg \max_{\eta} \mathbb{E}_{\phi} \mathbb{E}_{\theta_0 \sim p(\theta_0)}[V_{\phi, \theta_0}(\mathcal{F}_{\eta})]$ . The function  $V_{\phi}: \mathcal{H} \mapsto \mathbb{R}$  defines the expected return of the student policy on  $\phi$  at the end of the its lifetime (after  $N$  updates) using the update rule parameterized by  $\eta \in \mathcal{H}$ , when trained from an initialization  $\theta_0 \sim p(\theta_0)$ .
3.2 Approximating Minimax-Regret
As in the traditional UED setting, regret is generally intractable to compute. A range of scoring functions have been proposed to approximate regret, often deriving the approximation from value loss [Jiang et al., 2021b,a]. Two consistently high performing metrics are L1 value loss and positive value loss, which are equal to the episodic mean of absolute and positive GAE [Schulman et al., 2016] terms respectively.
In order to exploit the structure of our problem setting, we propose an alternative approximation which we refer to as algorithmic regret (AR). In this, we co-train an ant agonist agent using a manually designed RL algorithm  $\mathcal{A}$  (e.g., A2C [Mnih et al., 2016], PPO [Schulman et al., 2017]) in parallel to the protagonist agent trained by GROOVE. AR is then computed from the difference in final performance against the antagonist,
$$
\operatorname {R e g r e t} ^ {\mathcal {A}} (\eta , \phi) = V _ {\phi} (\mathcal {A}) - V _ {\phi} (\eta) \tag {4}
$$
where  $V_{\phi}(\mathcal{A})$  denotes the expected return of the antagonist policy when trained with  $\mathcal{A}$ .
We evaluate AR against both L1 value loss and positive value loss in Section 4.4, demonstrating improved generalization performance on Min-Atar.
3.3 Environment Design
Following the dual-curriculum design paradigm from Jiang et al. [2021a], a large class of UED methods can be represented as a combination of two teachers: a curator and a generator. Here, the level generator is a generative model that is optimized to produce regret-maximizing levels, whilst the level curator maintains a set of previously-visited high-regret levels to be replayed by the agent. The generator provides a slowly adapting mechanism for environment design, allowing the method to design new levels without random sampling, whilst the curator provides a quickly-adapting replay buffer of useful levels.
In PMO, a single sample of environment regret requires an agent to be trained to convergence and subsequently evaluated on that environment. This is significantly less sample efficient than traditional UED, where regret is measured from a single rollout of the current policy. Furthermore, generator-based methods for UED (e.g. PAIRED) have been shown to achieve lower performance and sample efficiency than curation-based approaches [Jiang et al., 2021a, Parker-Holder et al., 2022a]. Due to this, we design GROOVE using PLR (Section 2.3), which curates randomly generated levels, avoiding the need to train a level generator.
The meta-training loop for GROOVE is presented in Algorithm 1 and Figure 2.
Algorithm 1 GROOVE meta-training
input: Environment set  $\Phi$  agent parameter initialization function  $p(\theta)$    
initialize: Meta-parameters  $\eta$  PLR level buffer  $\Lambda$  ,agent-environment lifetimes  $\{\theta ,\phi \} _i$  repeat   
for all lifetimes  $\{\theta ,\phi \} _i$  do for update  $\leftarrow 1$  to  $K$  do Rollout agent  $\pi_{\theta}$  on environment  $\phi$  Update  $\theta$  with  $\mathcal{F}_{\eta}$  (Equation 2) end for Compute meta-gradient for  $\eta$  with updated  $\theta$  (Equation 1) if lifetime over then Evaluate regret approximation  $\mathcal{R}\gets \mathrm{REGRET}^{*}(\eta ,\phi)$  with final  $\theta$  Update level buffer  $\Lambda$  with  $\mathcal{R}$  and  $\phi$  Reinitialize lifetime  $(\theta ,\phi)\sim p(\theta)\times \Lambda$  end if end for Update  $\eta$  with accumulated meta-gradients   
until  $\eta$  converges
Our experiments are designed to determine (1) how the meta-training distribution impacts OOD generalization in PMO, (2) how well AR identifies informative levels for generalization, and (3) the effectiveness of GROOVE at generating curricula for generalization using this metric. To achieve this, we first manually examine how informative and diverse meta-training distributions improve the generalization performance of LPG (Section 4.2). We then evaluate curricula generated by AR against those generated by randomly sampling and handcrafting (Section 4.3), demonstrating the ability to identify informative levels. Following this, we evaluate GROOVE against LPG (Section 4.4), demonstrating the impact of environment design for improving both in-distribution robustness and generalization performance. Finally, we evaluate GROOVE with AR against baseline metrics from the UED literature (Section 4.5), showing only AR consistently improves generalization.
Training Environment For meta-training, we use a generalization of the tabular Grid-World environment presented by Oh et al. [2020]. In this environment distribution, a task is specified by the maximum episode length, grid size, wall placement, start position, and number of objects, whilst the objects themselves vary in position, reward, and probabilities of respawning or terminating the episode. This space contains tasks encapsulating thematic challenges in RL, including exploration, credit assignment, and stochasticity.
However, these challenges are notably sparse over the environment distribution. For instance, maze-like arrangements of walls induce a hard exploration challenge by creating anomalously long shortest-path lengths to objects, but are rare under uniform sampling. This captures the need for environment design, in order to discover complex and informative structures required for generalization.
Testing Environments The purpose of our evaluation is to determine the generalization performance of the algorithms we consider, i.e., the expected return on real-world RL tasks. In order to approximate this, we evaluate on Atari [Bellemare et al., 2013], an archetypal RL benchmark, as well as its simplified counterpart Min-Atar [Young and Tian, 2019] for our intermediate results.
Model Architecture and Implementation For our learned optimizer, we use the model architecture proposed in LPG. Since GROOVE is agnostic to the underlying meta-optimization method, we select LPG due to its state-of-the-art generalization performance on unseen tasks, in addition to the prior analysis of training distribution performed on LPG, which we build upon in Section 4.2. Further comparison to prior meta-optimization methods is presented in Section 5. Our experiments were executed on two to five servers, containing eight GPUs each (ranging in performance from 1080-Ti to V100). Model hyperparameters can be found in the supplementary materials and the project repository is available at https://github.com/EmptyJackson/groove.
4.2 Designing Meta-Training Distributions for Generalization
Unlike in standard RL, the impact of meta-training distributions on OOD generalization in PMO has not been explored in depth. One attempt at this came from Oh et al. [2020], who evaluate the generalization performance of LPG after meta-training on three different environment sets, varying both the number of tasks and the environments that the tasks are sampled from. By demonstrating improved transfer to Atari, the authors claim that two factors improve generalization performance:
1. Task diversity (the number of training tasks), and
2. How informative the tasks are for generalization.2
Whilst these results suggest a relationship between the meta-training distribution and generalization performance, the strength of these claims is limited by the number of environment sets (three) and confounding of these factors.
Quote: "specific types of training environments... improved generalization performance." We refer to these types of environments as being informative for generalization.
To analyze the impact of task diversity, we sample Grid-World subsets of various sizes, before meta-training LPG on each of these and evaluating on Min-Atar (Figure 3). By generating environment sets with i.i.d. sampling from the Grid-World environment distribution, we remove the informativeness of tasks as a confounding factor. We observe a significant  $(p < 0.05)$  positive correlation in performance with number of levels in the aggregate task return, supporting the first claim. Furthermore, when considering the per-task breakdown of results (see supplementary materials), we observe a significant positive correlation on three of the four tasks, with the remaining task having a weak positive correlation, thereby supporting the first claim.
We investigate the second factor using a set of handcrafted Grid-World configurations proposed by Oh et al. [2020] (see supplementary materials). These are manually designed to emphasize stochasticity and credit assignment, two
key challenges in RL, making them more informative for generalization than randomly sampled Grid-World configurations. At the same number of levels, we observe an improvement in performance from meta-training on handcrafted levels against random levels. Moreover, training on five handcrafted levels exceeds the performance of  $2^{6} = 64$  random levels, demonstrating the need for task distributions to contain tasks that are both informative and diverse.