Grounding Aleatoric Uncertainty in Unsupervised Environment Design

1 INTRODUCTION . Adaptive curricula , which dynamically adjust the distribution of training environments to optimize the performance of the resulting policy , have played a key role in many recent achievements in deep reinforcement learning ( RL ) . Applications have spanned both single-agent RL ( Portelas et al. , 2020 ; Wang et al. , 2019 ; Zhong et al. , 2020 ; Justesen et al. , 2018 ) , where adaptation occurs over environment variations , and multi-agent RL ( MARL ) , where adaptation can additionally occur over co-players ( Silver et al. , 2016 ; Vinyals et al. , 2019 ; Stooke et al. , 2021 ) . By presenting the agent with challenges at the threshold of its abilities , such methods demonstrably improve the sample efficiency and the generality of the final policy ( Matiisen et al. , 2017 ; Dennis et al. , 2020 ; Jiang et al. , 2021b ; a ) . This work introduces a fundamental problem relevant to adaptive curriculum learning methods for RL , which we call curriculum-induced covariate shift ( CICS ) . Analogous to the covariate shift that occurs in supervised learning ( SL ) , CICS refers to a mismatch between the input distribution at training and test time , and in this case , specifically when the distribution shift is caused by the selective sampling performed by an adaptive curriculum . While there may be cases in which CICS impacts model performance in SL , adaptive curricula for SL have generally not been found to be as impactful as in RL ( Wu et al. , 2021 ) . Therefore , here we focus on addressing this problem specifically as it arises in the RL setting , and leave investigation of its potential impact in SL to future work . To establish precise language around adaptive curricula , we cast our discussion under the lens of Unsupervised Environment Design ( UED , Dennis et al. , 2020 ) . UED provides a formal problem description for which curriculum learning is the solution , by defining the Underspecified POMDP ( UPOMDP ; see Section 2 ) , which expands the classic POMDP with a set of free parameters Θ , representing the dimensions along which the environment may vary across episodes . The goal of UED is then to adapt distributions over Θ , so to maximize some objective , which could be tied to an RL agent ’ s performance over this distribution . This allows us to view adaptive curricula as emerging via a multi-agent game between a teacher that proposes environments with parameters θ ∼ Θ and a student that learns to solve them . In addition to the notational clarity it provides , this formalism lends the analysis of adaptive curricula to useful game theoretic constructs , such as Nash equilibria ( NE , Nash et al. , 1950 ) . This game-theoretic view has led to the development of curriculum methods with principled robustness guarantees , such as PAIRED ( Dennis et al. , 2020 ) and Prioritized Level Replay ( PLR , Jiang et al. , 2021a ) , which showed that curricula optimizing for the student ’ s regret lead to minimax regret ( Savage , 1951 ) policies at Nash equilibria , implying the agent can solve all solvable environments within the training domain . While other methods can also be cast in this framework , they do not hold the same desirable property at equilibrium . For this reason , we will focus on addressing CICS for regret-maximizing UED , but note that our solution can be used with other UED methods . To see how the CICS can be problematic , consider the simplified case of training a self-driving car in simulation , so it learns to take the fastest route from home to office . Suppose traffic data shows on 70 % of the days , Route 1 is faster than Route 2 . Moreover , on any given day the self-driving car can not infer which route is faster ahead of time , so always picking Route 1 is faster in expectation . To support training a policy using an adaptive curriculum , one could build a simulator which sets road conditions per episode based on a random day sampled from the traffic data . However , adaptive curriculum over the traffic settings may oversample days when Route 1 is closed—perhaps because it finds the agent needs more practice on Route 2—shifting the best choice to Route 2 in training . In fact , methods for minimax regret UED like PLR would keep shifting the distribution of fastest route , to maximize the agent ’ s regret . Their curriculum dynamics would map to the zero-sum game of matching pennies , in which one player wins for guessing whether the other chose heads or tails ; the NE corresponds to each player randomly playing each option half the time . Randomly picking the route in this way is suboptimal , because in reality Route 1 is faster in expectation . This example is depicted in Figure 1 , where the two routes are replaced by an apple and banana . If , on a given day , the faster route could be identified before having to pick one , the agent could choose optimally . Instead , it is an aleatoric parameter , inducing irreducible uncertainty in the limit of infinite experiential data ( Der Kiureghian & Ditlevsen , 2009 ) . When CICS occurs over such parameters , with respect to a ground-truth distribution of environments P ( θ ) , the learned policy can be suboptimal with respect to P . It can therefore be useful to ground—that is , to constrain—the aleatoric parameters Θ′ ⊂ Θ to P ( θ′ ) when it is known or can be learned , as in simulation or from real-world data . However , grounding all θ is undesirable , preventing the curriculum from sampling enough opportunities to learn from useful scenarios with low support under P ( θ ) . In this work , we formalize the problem of CICS in RL , and provide a solution by proposing a UED method to find robustly Bayes-optimal policies where θ′ is grounded to P ( θ′ ) . Our solution called Sample-Matched PLR ( SAMPLR ) extends PLR , a state-of-the-art UED algorithm , by constraining the advantage estimates to match those observed if training under P ( θ′ ) . This advantage correction adapts Off-Belief Learning ( Hu et al. , 2021 ) from cooperative MARL , revealing an intriguing connection between curriculum biases observed in single and multi-agent RL . Our experiments in challenging environments based on the NetHack Learning Environment ( NLE , Küttler et al. , 2020 ) demonstrate that SAMPLR learns near-optimal policies under CICS , in cases where standard PLR fails . 2 BACKGROUND . 2.1 UNSUPERVISED ENVIRONMENT DESIGN . The problem of Unsupervised Environment Design ( UED , Dennis et al . ( 2020 ) ) is the problem of automatically generating an adaptive distribution of environments which will lead to policies that successfully transfer within a target domain . The domain of possible environments is represented by an Underspecified POMDP ( UPOMDP ) , which adds a set of free parameters to the standard definition of a POMDP , along which each concrete instantiation , or level , of the UPOMDP . For instance , these free parameters can be the position of obstacles in a maze , or friction coefficients in a physics-based task . Formally a UPOMDP is defined as a tupleM = 〈A , O , Θ , S , T , I , R , γ〉 , where A is a set of actions , O is a set of observations , Θ is a set of free parameters , S is a set of states , T : S ×A×Θ→∆ ( S ) is a transition function , I : S → O is an observation ( or inspection ) function , R : S → R is a reward function , and γ is a discount factor . UED typically approaches the curriculum design problem as training a teacher agent that co-evolves an adversarial curriculum for a student agent , for example , by maximizing the student ’ s regret . We will focus on a recent UED algorithm called Robust Prioritized Level Replay ( PLR⊥ , Jiang et al. , 2021b ) , which performs environment design via random search . PLR maintains a buffer of the most useful levels for training , according to some learning potential score—typically based on a regret approximation , such as the positive value loss—and with probability p , actively samples the next training level from this level buffer instead of the ground-truth training distribution . This selective-sampling mechanism has been demonstrated to greatly improve sample-efficiency and generalization in several domains , while provably leading to a minimax regret policy for the student at NE . In maximizing regret , PLR curricula naturally avoid unsolvable levels , which have no regret . 2.2 OFF-BELIEF LEARNING . In cooperative MARL , self-play promotes the formation of cryptic conventions—arbitrary sequences of actions that allow agents to communicate information about the environment state . These conventions are learned jointly among all agents during training , but are arbitrary and hence , indecipherable to independently-trained agents or humans at test time . Crucially , this leads to policies that fail to perform zero-shot coordination ( ZSC , Hu et al. , 2020 ) , where independentlytrained agents must cooperate successfully without additional learning steps , or ad-hoc team play . Off-Belief Learning ( OBL ) resolves this problem by forcing agents to assume their co-players act according to a fixed , known policy π0 until the current time t , and optimally afterwards , conditioned on this assumption . If π0 is playing uniformly random , this removes the possibility of forming arbitrary conventions . Formally , let G be a decentralized , partially-observable MDP ( Dec-POMDP , Bernstein et al. , 2002 ) , with state s , joint action a , observation function Ii ( s ) for each player i , and transition function T ( s , a ) . Let the historical trajectory τ = ( s1 , a1 , ... at−1 , st ) , and the action-observation history ( AOH ) for agent i be τ i = ( Ii ( s1 ) , a1 , ... , at−1 , Ii ( st ) ) . Further , let π0 be an arbitrary policy , such as a uniformly random policy , and Bπ0 ( τ |τ i ) = P ( τt|τ it , π0 ) , a belief model predicting the current state , conditioned on the AOH of agent i and the assumption of co-players playing policy π0 until the current time t , and optimally according to π1 from t and beyond . OBL aims to find the policy π1 with the optimal , counter-factual value function , V π0→π1 ( τ i ) = Eτ∼Bπ0 ( τ i ) [ V π1 ( τ ) ] . ( 1 ) As the agent conditions its policy on the realized AOH τ i , while transition dynamics are based on states sampled from B , this mechanism is called a fictitious transition . In Section 5 , we show how OBL ’ s fictitious transition can be adapted to the single-agent curriculum learning setting to address CICS , by interpreting the curriculum designer in UED as a co-player . 3 RELATED WORK . The mismatch between training and testing distributions of input features is referred to as covariate shift , and has long served as a fundamental problem for the machine learning community . Covariate shifts have been extensively studied in supervised learning ( Vapnik & Chervonenkis , 1971 ; Huang et al. , 2006 ; Bickel et al. , 2009 ; Arjovsky et al. , 2019 ) . In RL , prior works have largely focused on covariate shifts due to training on off-policy data ( Sutton et al. , 2016 ; Rowland et al. , 2020 ; Espeholt et al. , 2018 ; Hallak & Mannor , 2017 ; Gelada & Bellemare , 2019 ; Thomas & Brunskill , 2016 ) including the important case of learning from demonstrations ( Pomerleau , 1988 ; Ross & Bagnell , 2010 ) . Recent work also aims to learn invariant representations robust to covariate shifts ( Zhang et al. , 2019 ; 2021 ) . More generally , CICS can be interpreted as a kind of sample-selection bias ( Heckman , 1979 ) . We believe this work to be the first to formalize and provide a solution to the problem of covariate shifts in reinforcement learning due to curriculum learning . Our method fixes a critical flaw that can cause curricula to fail under CICS—an important problem as curricula have been shown to be essential for training RL agents across many of the most challenging domains , including combinatorial gridworlds ( Zhong et al. , 2020 ) , Go ( Silver et al. , 2016 ) , StarCraft 2 ( Vinyals et al. , 2019 ) , and achieving comprehensive task mastery in open-ended environments ( Stooke et al. , 2021 ) . While this work focuses on PLR , other recent methods include minimax adversarial curricula ( Wang et al. , 2019 ; 2020 ) and curricula based on changes in return ( Matiisen et al. , 2017 ; Portelas et al. , 2020 ) . Most similar to our work , OFFER ( Ciosek & Whiteson , 2017 ) adapts a curriculum over transition functions and uses importance sampling to correct for biased gradient estimates . Unlike this work , Ciosek & Whiteson ( 2017 ) requires whitebox access to the transition function and does not directly study the impact of CICS on the learning dynamics . Curriculum methods have also been studied in goal-conditioned RL ( Florensa et al. , 2018 ; Campero et al. , 2021 ; Sukhbaatar et al. , 2018 ; OpenAI et al. , 2021 ) , though CICS does not occur here as goals are observed by the agent . Lastly , domain randomization ( DR , Sadeghi & Levine , 2017 ; Peng et al. , 2017 ) can be seen as a degenerate form of UED , though curriculum-based extensions of DR have also been studied ( Jakobi , 1997 ; Tobin et al. , 2017 ) . Prior work has also investigated methods for learning Bayes optimal policies under uncertainty about the task ( Zintgraf et al. , 2020 ; Osband et al. , 2013 ) , based on the framework of Bayes-adaptive MDPs ( BAMDPs ) ( Bellman , 1956 ; Duff , 2002 ) . In this setting , the agent can adapt to an unknown MDP over several episodes by acting to reduce its uncertainty about the identity of the MDP . In contrast , SAMPLR learns a robustly Bayes optimal policy for the case of zero-shot transfer . Further unlike these works , our setting assumes the distribution of certain aleatoric parameters are biased during training , which would lead to biased a posteriori uncertainty estimates with respect to the ground-truth distribution when optimizing for the BAMDP objective . Instead , SAMPLR proposes a means to correct for this bias assuming knowledge of the true environment parameters for each level , to which we can safely assume access in curriculum learning .

Summary: This work proposes a solution to the problem of covariate shift appeared in adaptive curriculum learning where the distribution of parameters of the environment at test time is different from the one at training time. This is caused in algorithms like PAIRED or PLR because the adaptive curriculum learning algorithm biases towards regret maximising environments and as such trains the learner on a biased distribution instead of the ground truth one. The authors propose a solution to the problem under the assumption that the designer knows the ground truth distribution of the parameters of the environment. The solution relies on using fictitious transitions to compute the value function, using Off Belief learning — an algorithm for cooperative MARL.