Grounding Aleatoric Uncertainty in Unsupervised Environment Design
Inflection AI
Edward Grefenstette
Abstract
Adaptive curricula in reinforcement learning (RL) have proven effective for producing policies robust to discrepancies between the train and test environment.
Recently, the Unsupervised Environment Design (UED) framework generalized RL curricula to generating sequences of entire environments, leading to new methods with robust minimax regret properties.
Problematically, in partially-observable or stochastic settings, optimal policies may depend on the ground-truth distribution over aleatoric parameters of the environment in the intended deployment setting, while curriculum learning necessarily shifts the training distribution.
We formalize this phenomenon as curriculum-induced covariate shift (CICS), and describe how its occurrence in aleatoric parameters can lead to suboptimal policies.
Directly sampling these parameters from the ground-truth distribution avoids the issue, but thwarts curriculum learning.
We propose SAMPLR, a minimax regret UED method that optimizes the ground-truth utility function, even when the underlying training data is biased due to CICS.
We prove, and validate on challenging domains, that our approach preserves optimality under the ground-truth distribution, while promoting robustness across the full range of environment settings.
Introduction
Adaptive curricula, which dynamically adjust the distribution of training environments to best facilitate learning, have played a key role in many recent achievements in deep reinforcement learning (RL).
Applications have spanned both single-agent RL [32, 50, 55, 22], where adaptation occurs over environment variations and multi-agent RL, where adaptation can additionally occur over co-players [40, 49, 41].
These methods demonstrably improve the sample efficiency and robustness of the final policy [25, 8, 21, 20], e.g. by presenting the agent with challenges at the threshold of its abilities.
In this paper we introduce and address 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) [18], CICS refers to a mismatch between the input distribution at training and test time.
In the case of RL, we will show this becomes problematic when the shift occurs over the aleatoric parameters of the environment—those aspects of the environment holding irreducible uncertainty even in the limit of infinite experiential data [9].
While in some cases,
CICS may impact model performance in SL, adaptive curricula for SL have generally not been found to be as impactful as in RL [52].
Therefore, we focus on addressing CICS specifically as it arises in the RL setting, leaving investigation of its potential impact in SL to future work.
[Figure]
To establish precise language around adaptive curricula, we cast our discussion under the lens of
Unsupervised Environment Design [UED, 8].
UED provides a formal problem description for which an optimal curriculum 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 aspects of the environment that may vary.
UED then seeks to adapt distributions over Θ to maximize some objective, potentially tied to the agent’s performance.
UED allows us to view adaptive curricula as emerging via a multi-player game between a teacher that proposes environments with parameters θ ∼ P (Θ) and a student that learns to solve them.
In addition to notational clarity, this formalism enables using game theoretic constructs, such as Nash equilibria [NE, 26], to analyze curricula.
This game-theoretic view has led to the development of curriculum methods with principled robustness guarantees, such as PAIRED [8] and Robust Prioritized Level Replay [PLR ⊥ , 20], which aim to maximize a student’s regret and lead to minimax regret [37] policies at NE.
Thus, at NE, the student can solve all solvable environments within the training domain.
However, in their current form the
UED robustness guarantees are misleading: if the UED curriculum deviates from a ground-truth distribution P (Θ) of interest, i.e. the distribution at deployment, with respect to aleatoric parameters Θ ′ ⊂
Θ , the resulting policies may be suboptimal under the ground-truth distribution P .
For a concrete example of how CICS can be problematic, consider the case of training a self-driving car to navigate potentially icy roads, when icy conditions rarely occur under P .
When present, the ice is typically hard to spot in advance; thus, the aleatoric parameters Θ ′ correspond to whether each section of the road is icy.
A priori, a curriculum should selectively sample more challenging icy settings to facilitate the agent’s mastery over such conditions.
However, this approach risks producing an overly-pessimistic agent (i.e. one that assumes that ice is common), driving slowly even in fair weather.
Such a policy leads to inadequate performance on P , which features ice only rarely.
We can preserve optimality on P by grounding the policy —that is, ensuring that the agent acts optimally with respect to the ground-truth utility function for any action-observation history τ and the implied ground-truth posterior over Θ :
[Equation]
where the ground-truth utility conditioned on X , U ( π | X ) , is defined to be E τ,θ ∼ P ( θ | X ) [ (cid:80) ∞ t =
γ t r t ] , for rewards r t and a discount γ .
We can ground the policy by grounding the training distribution , which means constraining the training distribution of aleatoric parameters P (Θ ′ ) to match P (Θ ′ ) .
This is trivially accomplished by directly sampling θ ′ ∼ P (Θ ′ ) , which we call naive grounding .
Unfortunately, this approach makes many curricula infeasible by removing the ability to selectively sample environment settings over aleatoric parameters.
Applying this strategy to the self-driving agent may result in a policy that is optimal in expectation under P where there is rarely ice, but nevertheless fails to drive safely on ice.
We wish to maintain the ability to bias a training distribution, since it is required for curriculum learning, while ensuring the resulting decisions remain optimal in expectation under P .
This goal is captured by the following objective:
[Equation]
where D is the training distribution of τ .
Under naive grounding, D is equal to P ( τ ) and Equation 2 reduces to U ( π ) .
To overcome the limitations of naive grounding, we develop an approach that allows D to deviate from P ( τ ) , e.g. by prioritizing levels most useful for learning, but still grounds the policy by evaluating decisions following potentially biased training trajectories τ according to
U ( π | τ ) .
Figure 1 summarizes this approach, and contrasts it with an ungrounded adaptive curriculum.
In summary this work presents the following contributions: i)
We first formalize the problem of CICS in RL in Section 3. ii)
Then, we present SAMPLR, which extends PLR ⊥ , a state-of-the-art UED method, to preserve optimality on P while training under a usefully biased training distribution in
Section 4.
iii)
We prove in Section 5 that SAMPLR promotes Bayes-optimal policies that are robust over all environment settings θ ∼ P (Θ) .
iv)
Our experiments validate these conclusions in two challenging domains, where SAMPLR learns highly robust policies, while PLR ⊥ fails due to CICS.
Background 2.1 Unsupervised Environment Design
Unsupervised Environment Design [UED, 8] 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 environment settings is represented by an Underspecified POMDP (UPOMDP), which models each environment instantiation, or level , as a specific setting of the free parameters that control how the environment varies across instances.
Examples of free parameters are the position of walls in a maze or friction coefficients in a physics-based task.
Formally a UPOMDP is defined as a tuple M = ⟨ A, O, Θ , S , T , I , R , γ ⟩ , where A is the action space, O is the observation space, Θ is the set of free parameters, S is the state space, T : S ×
A × Θ → ∆ ( S ) is the transition function, I : S
→ O is the observation function, R :
S
→ R is the reward function, and γ is the discount factor.
UED typically approaches the curriculum design problem as training a teacher agent that co-evolves an adversarial curriculum for a student agent, e.g. by maximizing the student’s regret.
2.2 Prioritized Level Replay
We focus on a recent UED algorithm called Robust Prioritized Level Replay [PLR ⊥ , 21], which performs environment design via random search.
PLR ⊥ maintains a buffer of the most useful levels for training, according to an estimate of learning potential—typically based on regret, approximated by a function of the temporal-difference (TD) errors incurred on each level.
For each episode, with probability p , PLR ⊥ actively samples the next training level from this buffer, and otherwise evaluates regret on a new level θ ∼ P (Θ) without training.
This sampling mechanism provably leads to a minimax regret policy for the student at NE, and has been shown to improve sample-efficiency and generalization.
The resultant regret-maximizing curricula naturally avoid unsolvable levels, which have no regret.
We provide implementation details for PLR ⊥ in Appendix A.
Curriculum-Induced Covariate Shift
Since UED algorithms formulate curriculum learning as a multi-agent game between a teacher and a student agent, we can formalize when CICS becomes problematic by considering the equilibrium point of this game: Let Θ be the environment parameters controlled by UED, P (Θ) , their groundtruth distribution, and P (Θ) , their curriculum distribution at equilibrium.
We use τ t to refer to the joint action-observation history (AOH) of the student until time t (and simply τ when clear from context).
Letting V ( π | τ t ) denote the value function under the curriculum distribution P (Θ) , we characterize an instance of CICS over Θ as problematic if the optimal policy under P (Θ) differs from that under the ground-truth P (Θ) for some τ t , so that
[Equation]
π V ( π | τ t ) . The value function V ( π | τ t ) with respect to P (Θ) can be expressed as a marginalization over θ :
[Equation]
Here, the notation P ( θ ) means P (Θ = θ ) , and the tilde on the ˜ P and ˜ V terms indicates independence from any distribution over Θ , as they both condition on θ .
Importantly, the value function under the curriculum distribution V ( π | τ t ) corresponds to Equation 3 with P replaced by P .
We see that
V ( π | τ t ) is unchanged for a given τ t when P ( θ ) is replaced with P ( θ ) if 1) P ( θ ∗ | τ t )
= 1 for some θ ∗ , and 2) P shares support with P .
Then ˜ P ( τ t | θ )
= 1 iff θ = θ ∗ and zero elsewhere.
In this case, the sums reduce to V ( π | τ t , θ ∗ ) , regardless of changing the ground-truth distribution P to P .
In other words, when Θ is fully determined given the current history τ , covariate shifts over Θ with respect to
P (Θ) have no impact on policy evaluation and thus the value function for the optimal policy.
If the first condition does not hold, the uncertainty over the value of some subset Θ ′ ⊂
Θ is irreducible given τ , making Θ ′ aleatoric parameters for the history τ .
Thus, assuming the curriculum shares support with the ground-truth distribution, covariate shifts only alter the optimal policy at τ when they occur over aleatoric parameters given τ .
Such parameters can arise when the environment is inherently stochastic or when the cost of reducing uncertainty is high.
Crucially, our analysis assumes P and P share support over Θ .
When this assumption is broken, the policy trained under the curriculum can be suboptimal for environment settings θ , for which
P ( θ ) = 0 and P ( θ ) > 0 .
In this paper, we specifically assume that P and P share support and focus on addressing suboptimality under the ground-truth P due to CICS over the aleatoric parameters Θ ′ .
This discussion thus makes clear that problematic CICS can be resolved by grounding the training distribution , i.e. enforcing the constraint P (Θ ′ | τ )
=
P (Θ ′ | τ ) for the aleatoric parameters of the environment.
This constraint results in grounding the policy , i.e. ensuring it is optimal with respect to the ground-truth utility function based on P (Equation 1).
As discussed, naive grounding satisfies this constraint by directly sampling θ ′ ∼ P (Θ ′ )
, at the cost of curricula over Θ ′ .
This work develops an alternative for satisfying this constraint while admitting curricula over Θ ′ .
Sample-Matched PLR (SAMPLR)
Randomly initialize policy π ( ϕ ) , an empty level buffer Λ of size K , and belief model B ( s t | τ ) . while not converged do
Sample replay-decision Bernoulli, d ∼ P D ( d ) if d = 0 or | Λ | = 0 then
Sample level θ from level generator Collect π ’s trajectory τ on θ , with a stop-gradient ϕ ⊥ else
Use PLR to sample a replay level from the level store, θ ∼ Λ
Collect fictitious trajectory τ ′ on θ , based on s ′ t ∼ B
Update π with rewards R ( τ ′ ) end
Compute PLR score, S = score ( τ ′ , π )
Update Λ with θ using score S end
We now describe a general strategy for addressing CICS, and apply it to
PLR ⊥ , resulting in Sample-Matched PLR (SAMPLR).
This new UED method features the robustness properties of PLR ⊥ while mitigating the potentially harmful effects of
CICS over the aleatoric parameters Θ ′ .
As discussed in Section 3, CICS become problematic when the covariate shift occurs over some aleatoric subset Θ ′ of the environment parameters Θ , such that the expectation over Θ ′ influences the optimal policy.
Adaptive curriculum methods like
PLR ⊥ prioritize sampling of environment settings where the agent experiences the most learning.
While such a curriculum lets the agent focus on correcting its largest errors, the curriculum typically changes the distribution over aleatoric parameters Θ ′ , inducing bias in the resulting decisions.
Ideally, we can eliminate this bias, ensuring the resulting policy makes optimal decisions with respect to the ground-truth utility function, conditioned on the current trajectory: [Equation]
[Equation]
A naive solution for grounding is to simply exclude Θ ′ from the set of environment parameters under curriculum control.
That is, for each environment setting proposed by the curriculum, we resample θ ′ ∼ P .
We refer to this approach as naive grounding .
Naive grounding forces the expected reward and next state under each transition at the current AOH τ to match that under P .
Thus, optimal policies under naive grounding must be optimal with respect to the ground-truth distribution over θ ′ .
While technically simple, naive grounding suffers from lack of control over Θ ′ .
This limitation is of no concern when the value of Θ ′ does not alter the distribution of τ until the terminal transition, e.g. when Θ ′ is the correct choice in a binary choice task, thereby only influencing the final, sparse reward when the right choice is made.
In fact, our initial experiment in Section 6 shows naive grounding performs well in such cases.
However, when the value of Θ ′ changes the distribution of τ before the terminal transition, the agent may benefit from a curriculum that actively samples levels which promote learning robust behaviors under unlikely events.
Enabling the full benefits of the curriculum in such cases requires the curriculum to selectively sample values of Θ ′ .
Instead of naive grounding, we aim to ground only the policy updates, allowing the curriculum to bias the training distribution.
This can be accomplished by optimizing the following objective:
[Equation]
To achieve this, we replace the reward r t and next state s t +1 with counterfactual values that would be experienced if θ ′ were consistent with τ and P , so that θ ′ ∼ P ( θ ′ | τ ) .
This substitution occurs by simulating a fictitious transition , where the fictitious state is sampled as s ′ t ∼ B ( s ′ t | τ ) , the action as a t ∼ π ( ·| τ ) (as per usual), the fictitious next state as s ′ t [-+
-]+1 =
T ( s ′ t , a t ) , and the fictitious reward as r ′ t =
R ( s ′ t +1 ) .
The belief model B ( s ′ t | τ ) is the ground-truth posterior of the current state given τ :
[Equation]
Fictitious transitions, summarized in Figure 2, ground the observed rewards and state transitions to P .
Should training on these transitions lead to an optimal policy over Θ , this policy will also be optimal with respect to P .
We prove this property in Section 5.
Fictitious transitions thus provide the benefit of naive grounding without giving up curriculum control over Θ ′ .
In general, we implement B as follows:
Given P (Θ ′ ) as a prior, we model the posterior P ( θ ′ | τ ) with Bayesian inference.
The posterior could be learned via supervised learning with trajectories collected from the environment for a representative selection of θ ′ .
Further, we may only have limited access to P (Θ) throughout training, for example, if sampling P (Θ) is costly.
In this case, we can learn an estimate ˆ P (Θ ′ ) from samples we do collect from P (Θ) , which can occur online.
We can then use ˆ P (Θ ′ ) to inform the belief model.
SAMPLR, summarized in Algorithm 1, incorporates this fictitious transition into PLR ⊥ by replacing the transitions experienced in replay levels sampled by PLR ⊥ with their fictitious counterparts, as PLR ⊥ only trains on these trajectories.
PLR ⊥ uses PPO with the Generalized Advantage Estimator [GAE, 38] as the base RL algorithm, where both advantage estimates and value losses can be written in terms of one-step TD errors δ t at time t .
Training on fictitious transitions then amounts to computing these TD errors with fictitious states and rewards: δ t = r ′ t + V ( s ′ t ) − V ( s ′ t +1 ) .
Importantly, because PLR ⊥ provably leads to policies that minimize worst-case regret over all θ at
NE, SAMPLR enjoys the same property for θ ∼ P (Θ) .
A proof of this fact is provided in Section 5.
Applying SAMPLR requires two key assumptions:
First, the simulator can be reset to a specific state, which is often true, as RL largely occurs in resettable simulators or those that can be made to do so.
When a resettable simulator is not available, a possible solution is to learn a model of the environment which we leave for future work.
Second, we have knowledge of P (Θ ′ ) .
Often, we know P a priori , e.g. via empirical data or as part of the domain specification, as in games of chance.
The Grounded Optimality of SAMPLR
Training on fictitious transitions is a method for learning an optimal policy with respect to the groundtruth utility function U D ( π ) over the distribution D of training trajectories τ , defined in Equation 5.
When D corresponds to the distribution of trajectories on levels θ ∼ P (Θ) , U D ( π ) reduces to the ground-truth utility function, U ( π ) .
For any UED method, our approach ensures that, in equilibrium, the resulting policy is Bayes-optimal with respect to P (Θ) for all trajectories in the support of D .
Remark 1.
If π ∗ is optimal with respect to the ground-truth utility function U D ( π ) then it is optimal with respect to the ground-truth distribution P (Θ) of environment parameters on the support of D .
Proof.
By definition we have π ∗ ∈ arg max π ∈ Π { U D ( π ) }
= arg max π ∈ Π
{ E τ ∼D (cid:2) U
( π | τ ) (cid:3) } .
Since π can condition on the initial trajectory τ , the action selected after each trajectory can be independently optimized.
Therefore, for all τ ∈ D , π ∗ ∈ arg max π ∈ Π { U ( π | τ ) } implying that π ∗ is the optimal policy maximizing U ( π | τ ) .
Thus, assuming the base RL algorithm finds Bayes-optimal policies, a UED method that optimizes the ground-truth utility function, as done by SAMPLR, results in Bayes-optimal performance over the ground-truth distribution.
If the UED method maximizes worst-case regret, we can prove an even stronger property we call robust ϵ -Bayes optimality .
Let U θ ( π ) be the ground-truth utility function for π on the distribution D πθ of initial trajectories sampled from level θ , so that U θ ( π ) =
U D πθ ( π ) .
Given a policy π maximizing U θ ( π ) , we say that π is robustly ϵ -Bayes optimal iff for all θ in the domain of P (Θ) and all π ′ , we have
[Equation]
Note how this property differs from being simply ϵ -Bayes optimal, which would only imply that
[Equation]
Robust ϵ -Bayes optimality requires π to be ϵ -optimal on all levels θ in the support of the groundtruth distribution, even those rarely sampled under P (Θ) .
We will show that at ϵ -Nash equilibrium,
SAMPLR results in a robustly ϵ -Bayes optimal policy for the ground-truth utility function U θ ( π ) .
In contrast, training directly on levels θ ∼ P (Θ) results in a policy that is only ϵ -Bayes optimal.
Theorem 1.
If π ∗ is ϵ -Bayes optimal with respect to U (cid:98) D ( π ) for the distribution (cid:98) D of trajectories sampled under π over levels maximizing the worst-case regret of π , as occurs under SAMPLR, then π ∗ is robustly ϵ -Bayes optimal with respect to the ground-truth utility function, U ( π ) .
Proof.
Let π ∗ be ϵ -optimal with respect to U (cid:98) D ( π ) where (cid:98) D is the trajectory distribution under π on levels maximizing the worst-case regret of π .
Let π ∗ be an optimal grounded policy.
Then for any θ ,
[Equation]
The first inequality follows from (cid:98) D being trajectories from levels that maximize worst-case regret with respect to π ∗ , and the second follows from π ∗ being ϵ -optimal on U (cid:98) D ( π ) .
Rearranging terms gives the desired condition.
Experiments
Our experiments first focus on a discrete, stochastic binary choice task, with which we validate our theoretical conclusions by demonstrating that CICS can indeed lead to suboptimal policies.
Moreover, we show that naive grounding suffices for learning robustly optimal policies in this setting.
However, as we have argued, naive grounding gives up control of the aleatoric parameters Θ ′ and thus lacks the ability to actively sample scenarios helpful for learning robust behaviors—especially important when such scenarios are infrequent under the ground-truth distribution P (Θ) .
SAMPLR induces potentially biased curricula, but retains optimality under P (Θ) by matching transitions under P (Θ ′ ) with those under P (Θ ′ ) .
We assess the effectiveness of this approach in our second experimental domain, based on the introductory example of driving icy roads.
In this continuous-control driving domain, we seek to validate whether SAMPLR does in fact learn more robust policies that transfer to tail cases under P (Θ ′ ) , while retaining high expected performance on the whole distribution P (Θ ′ ) .
All agents are trained using PPO [39] with the best hyperparameters found via grid search using a set of validation levels.
We provide extended descriptions of both environments alongside the full details of our architecture and hyperparameter choices in Appendix C.
[Figure]
6.1 Stochastic Fruit Choice
We aim to demonstrate the phenomenon of CICS in Stochastic Fruit Choice, a binary choice task, where the aleatoric parameter determines the correct choice.
This task requires the agent to traverse up to eight rooms, and in the final room, decide to eat either the apple or banana.
The correct choice θ ′ is fixed for each level, but hidden from the agent.
Optimal decision-making depends on the ground-truth distribution over the correct fruit, P (Θ ′ ) .
This task benefits from a curriculum over the number of rooms, but a curriculum that selectively samples over both room layout and correct fruit choice can lead to suboptimal policies.
Figure 3 shows example levels from this environment.
This domain presents a hard exploration challenge for RL agents, requiring robust navigation across multiple rooms.
Further, this environment is built on top of MiniHack [36], enabling integration of select game dynamics from the NetHack Learning Environment [23], which the agent must master to succeed: To go from one room to the next, the agent needs to learn to kick the locked door until it opens.
Upon reaching the final room, the agent must then apply the eat action on the correct fruit.
Let π A be the policy that always chooses the apple, and π B , the banana.
If the probability that the goal is the apple is P ( A ) =
q , then the expected return is R A q under π A and R B (1 − q ) under π B .
The optimal policy is π A when q > R B / ( R A + R B ) , and π B otherwise.
Domain randomization (DR), which directly samples each level θ ∼ P ( θ ) , optimizes for the correct ground-truth P (Θ ′ ) , but will predictably struggle to solve the exploration challenge.
PLR ⊥ may induce curricula easing the exploration problem, but can be expected make the correct fruit choice oscillate throughout training to maximize regret, leading to problematic CICS.
We set R A = 3 , R B = 10 , and q = 0 .
7 , making π B optimal with an expected return of 3 . 0 .
We compare the train and test performance of agents trained with DR, PLR ⊥ , and PLR ⊥ with naive grounding over 200M training steps in Figure 3.
In this domain, SAMPLR reduces to naive grounding, as θ ′ only effects the reward of a terminal transition, making fictitious transitions equivalent to real transitions for all intermediate time steps.
We see that DR struggles to learn an effective policy, plateauing at a mean return around 1 . 0 , while PLR ⊥ performs the worst.
Figure 6 in Appendix B shows that the PLR ⊥ curriculum exhibits much higher variance in q , rapidly switching the optimal choice of fruit to satisfy its regret-maximizing incentive, making learning more difficult.
In contrast,
PLR ⊥ with naive grounding constrains q = 0 .
7 , while still exploiting a curriculum over an increasing number of rooms, as visible in Figure 6.
This grounded curriculum results in a policy that solves more complex room layouts at test time.
Figures 5 and 6 in Appendix B additionally show how the
SAMPLR agent’s choices converge to π B and how the size of SAMPLR’s improvement varies under alternative choices of q in { 0 . 5 , 0 . 3 } .
6.2 Zero-Shot Driving Formula 1 Tracks with Black Ice
We now turn to a domain where the aleatoric parameters influence the distribution of τ t at each t , thereby creating opportunities for a curriculum to actively sample specific θ ′ to promote learning on biased distributions of τ t .
We base this domain on the black ice driving scenario from the introduction of this paper, by modifying the CarRacingBezier environment in [20].
In our version, each track tile has black ice with probability q , in which case its friction coefficient is 0, making acceleration and
[Figure]
e t
[Figure]
r n
[Figure]
Naive grounding
SAMPLR braking impossible.
This task is especially difficult, since the agent cannot see black ice in its pixel observations.
Figure 4 shows example tracks with ice rendered for illustration purposes.
The episodic returns scale linearly with how much of the track is driven and how quickly this is accomplished.
As success requires learning to navigate the challenging dynamics over ice patches [Table] a curriculum targeting more difficult ice configurations should lead to policies more robust to black ice.
Here, the ground-truth distribution P (Θ ′ ) models the realistic assumption that most days see little to no ice.
We therefore model the probability of ice per tile as q ∼ Beta ( α, β ) , where α = 1 , β = [Table]
[Table]
We test the hypothesis that SAMPLR’s regret-maximizing curriculum results in policies that preserve optimal performance on the ground-truth distribution P (Θ ′ ) , while being more robust to tail cases compared to DR and PLR ⊥ with naive grounding.
We expect standard PLR ⊥ to underperform all methods due to CICS, leading to policies that are either too pessimistic or too optimistic with respect to the amount of ice.
These baselines provide the controls needed to distinguish performance changes due to the two grounding approaches and those due to the underlying curriculum learning method.
We train agents with each method for 5M and test zero-shot generalization performance on the
Formula 1 (F1) tracks from the CarRacingF1 benchmark, extended to allow each track segment to have black ice with probability q in { 0 . 0 , 0 . 2 , 0 . 4 , 0 . 6 , 0 . 8 } .
These test tracks are significantly longer and more complex than those seen at training, as well as having a higher rate of black ice.
To implement SAMPLR’s belief model, we use a second simulator as a perfect model of the environment.
At each time step, this second simulator, which we refer to as the fictitious simulator , resets to the exact physics state of the primary simulator, and its icy tiles are resampled according to the exact posterior over the aleatoric parameter q = θ ′ , such that θ ′ ∼ P
θ ′ | τ ) , ensuring the future uncertainty is consistent with the past.
The agent decides on action a t based on the current real observation o t , and observes the fictitious return r ′ t and next state s ′ t +1 determined by the fictitious simulator after applying a t in state s ′ t ∼ P ( s ′ t | τ, θ ′ ) .
This dual simulator arrangement, fully detailed in Appendix A.2, allows us to measure the impact of training on fictitious transitions independently of the efficacy of a model-based RL approach.
Further, as the training environment in RL is most often simulation (e.g. in sim2real), this approach is widely applicable.
SAMPLR outperforms all baselines in zero-shot transfer to higher ice rates on the full F1 benchmark and attains a statistically significant improvement at p < 0 .
001 when transferring to q = 0 .
4 and q = 0 . 6 , and p < 0 .
05 when q = 0 . 8 .
Importantly, SAMPLR outperforms PLR ⊥ with naive grounding, indicating that SAMPLR exploits specific settings of Θ ′ to better robustify the agent against rare icy conditions in the tail of P (Θ ′ ) .
Indeed, Figure 4 shows that on average, SAMPLR exposes the agent to more ice per track tile driven, while PLR ⊥ underexposes the agent to ice compared to DR and naive grounding, suggesting that under PLR ⊥ agents attain higher regret on ice-free tracks—a likely outcome as ice-free tracks are easier to drive and lead to returns, with which regret scales.
Unfortunately, this results in PLR ⊥ being the worst out of all methods on the ground-truth distribution.
SAMPLR and naive grounding avoid this issue by explicitly matching transitions to those under P at τ .
As reported in Table 1, SAMPLR matches the baselines in mean performance across all F1 tracks under P (Θ ′ ) , indicating that despite actively sampling challenging θ ′ , it preserves performance under P (Θ ′ ) , i.e. the agent does not become overly cautious.
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 [48, 18, 5, 2].
In RL, prior works have largely focused on covariate shifts due to training on off-policy data [44, 34, 11, 14, 13, 46] including the important case of learning from demonstrations [31, 33].
Recent work also aimed to learn invariant representations robust to covariate shifts [53, 54].
More generally, CICS is a form of sample-selection bias [15].
Previous methods like OFFER [7] considered correcting biased transitions via importance sampling [43] when optimizing for expected return on a single environment setting, rather than robust policies over all environments settings.
We believe our work provides the first general formalization and solution strategy addressing curriculum-induced covariate shifts (CICS) for RL.
The importance of addressing CICS is highlighted by recent results showing curricula to be essential for training RL agents across many of the most challenging domains, including combinatorial gridworlds [55], Go [40], StarCraft 2 [49], and achieving comprehensive task mastery in open-ended environments [41].
While this work focuses on PLR ⊥ , other methods include minimax adversarial curricula [30, 50, 51] and curricula based on changes in return [25, 32].
Curriculum methods have also been studied in goal-conditioned RL [12, 6, 42, 27], though CICS does not occur here as goals are observed by the agent.
Lastly, domain randomization [DR, 35, 29] can be seen as a degenerate form of UED, and curriculum-based extensions of DR have also been studied [19, 47].
Prior work has also investigated methods for learning Bayes optimal policies under uncertainty about the task [56, 28], based on the framework of Bayes-adaptive MDPs (BAMDPs)
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 zero-shot transfer.
Further unlike these works, our setting assumes the distribution of some aleatoric parameters is biased during training, which would bias the 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, to which we can often safely assume access in curriculum learning.
Deeply related, Off-Belief Learning [OBL, 16] trains cooperative agents in self-play using fictitious transitions assuming all past actions of co-players follow a base policy, e.g. a uniformly random one.
Enforcing this assumption prevents agents from developing conventions that communicate private information to co-players via arbitrary action sequences.
Such conventions hinder coordination with independently trained agents or, importantly, humans.
SAMPLR can be viewed as adapting OBL to single-agent curriculum learning, where a co-player sets the environment parameters at the start of each episode (see Appendix D).
This connection highlights how single-agent curriculum learning is inherently a multi-agent problem, and thus problems afflicting multi-agent learning also surface in this setting; moreover, methods addressing such issues in one setting can then be adapted to the other.
Conclusion
This work characterized how curriculum-induced covariate shifts (CICS) over aleatoric environment parameters Θ ′ can lead to suboptimal policies under the ground-truth distribution over these parameters, P (Θ ′ ) .
We introduced a general strategy for correcting CICS, by training the agent on fictitious rewards and next states whose distribution is guaranteed to match what would be experienced under P (Θ ′ ) .
Our method SAMPLR augments PLR ⊥ with this correction.
By training on fictitious transitions, SAMPLR actively samples specific values of θ ′ that induce trajectories with greater learning potential, while still grounding the training data to P (Θ ′ ) .
Crucially, our experiments in challenging environments with aleatoric uncertainty showed that SAMPLR produces robust policies outperforming those trained with competing baselines that do not correct for CICS.
A core assumption made by SAMPLR and all other UED methods is the ability to reset the environment to arbitrary configurations of some set of free parameters.
While such resets can be difficult or impossible to perform in real world environments, in practice, this assumption is nearly always satisfied, as RL training largely occurs under a sim2real paradigm due to the additional costs of training in the wild.
Most RL simulators can either be directly reset to specific environment configurations or be straightforwardly made to do so.
SAMPLR thus provides a means to more fully exploit the affordances of a simulator to produce more robust policies: Policies trained with
SAMPLR retain optimality when transferred to the ground-truth distribution of aleatoric parameters in the real environment—a crucial property not satisfied by prior UED methods.
Importantly, the approach based on fictitious transitions used by SAMPLR can, in principle, be generally applied to prior UED methods to provide them with this desirable property.
Acknowledgments and Disclosure of Funding
We thank Robert Kirk and Akbir Khan for providing useful feedback on earlier drafts of this work.
Further, we are grateful to our anonymous reviewers for their valuable feedback.
MJ is supported by the FAIR PhD program at Meta AI.
This work was funded by Meta.