Distributionally Adaptive Meta Reinforcement Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
Meta-reinforcement learning algorithms provide a data-driven way to acquire policies that quickly adapt to many tasks with varying rewards or dynamics functions.
However, learned meta-policies are often effective only on the exact task distribution on which they were trained and struggle in the presence of distribution shiftof test-time rewards or transition dynamics.
In this work, we develop a framework for meta-RL algorithms that are able to behave appropriately under test-timedistribution shifts in the space of tasks.
Our framework centers on an adaptiveapproach to distributional robustness that trains a population of meta-policies tobe robust to varying levels of distribution shift.
When evaluated on a potentiallyshifted test-time distribution of tasks, this allows us to choose the meta-policy withthe most appropriate level of robustness, and use it to perform fast adaptation.
Weformally show how our framework allows for improved regret under distributionshift, and empirically show its efficacy on simulated robotics problems under awide range of distribution shifts.
Introduction
The diversity and dynamism of the real world require reinforcement learning (RL) agents thatcan quickly adapt and learn new behaviors when placed in novel situations.
Meta reinforcementlearning provides a framework for conferring this ability to RL agents, by learning a “meta-policy”trained to adapt as quickly as possible to tasks from a provided training distribution [35, 9, 30, 43].
Unfortunately, meta-RL agents are prone to overfitting to the distribution of tasks they are trainedon, and have been shown to behave erratically when asked to adapt to tasks beyond the trainingdistribution [4, 7].
As an example of this negative transfer, consider using meta-learning to teacha robot to navigate to goals quickly (illustrated in Figure 1).
The resulting meta-policy learns toquickly adapt and walk to any target location specified in the training distribution but explores poorlyand fails to adapt to any location not in that distribution.
Overfitting is particularly problematicfor the meta-learning setting, since the scenarios where we need the ability to learn quickly areusually exactly those where the agent experiences distribution shift.
This type of meta-distributionshift afflicts a number of real-world problems including autonomous vehicle driving [8], in-handmanipulation [14, 1], and quadruped locomotion [21, 19, 15], where the test-time task distributionmay not be well represented during training.
In this work, we study meta-RL algorithms that learn meta-policies resilient to task distribution shiftat test time.
One approach to enable this resiliency is to leverage the framework of distributionalrobustness [33], training meta-policies that prepare for distribution shifts by optimizing the worst-caseempirical risk against a set of task distributions which lie within a bounded distance from the originaltraining task distribution (often referred to as an uncertainty set )).
This allows meta-policies todeal with potential test-time task distribution shift, bounding their worst-case test-time regret for distributional shifts within the chosen uncertainty set.
However, choosing an appropriate uncertaintyset can be quite challenging without further information about the test environment, significantlyimpacting the test-time performance of algorithms under distribution shift.
Large uncertainty setsallow resiliency to a wider range of distribution shifts, but the resulting meta-policy adapts very slowlyat test time; smaller uncertainty sets enable faster test-time adaptation, but leave the meta-policybrittle to task distribution shifts.
Can we get the best of both worlds?
[Table]
<latexit sha1_
[Table]
="04ZwISNQzfwUIERNw0RDchdyb+M=">AAAB+nicbVBNS8NAEN34WetXqkcvwSJ4KkkV9Fj04rGC/YA2hM120i7dfLA7UUvMT/HiQRGv/hJv/hu3bQ7a+mDg8d4MM/
[Table]
/tm5aCt4lQyaLFYxLLrUwWCR9BCjgK6iQQa+
[Table]
6g5ilIUTIBFWq59gJuhmVyJmAvNxPFSSUjekQeppGNATlZrPTc+tEKwMriKWuCK2Z+
[Table]
[Figure]
ior
Our key insight is that we can prepare for a variety ofpotential test-time distribution shifts by constructing andtraining against different uncertainty sets at training time.
By preparing for adaptation against each of these uncertainty sets, an agent is able to adapt to a variety of potential test-time distribution shifts by adaptively choosing themost appropriate level of distributional robustness for thetest distribution at hand.
We introduce a conceptual framework called distributionally adaptive meta reinforcementlearning formalizing this idea.
At train time, the agentlearns robust meta-policies with widening uncertainty sets,preemptively accounting for different levels of test-timedistribution shift that may be encountered.
At test time,the agent infers the level of distribution shift it is facedwith, and then uses the corresponding meta-policy to adaptto the new task.
In doing so, the agent is able to adaptivelychoose the best level of robustness for the test-time taskdistribution, preserving the fast adaptation benefits of meta RL, while also ensuring good asymptotic performance under distribution shift.
We instantiate apractical algorithm in this framework (DiAMetR), using learned generative models to imagine newtask distributions close to the provided training tasks that can be used to train robust meta-policies.
The contribution of this paper is to propose a framework for making meta-reinforcement learningresilient to a variety of task distribution shifts, and DiAMetR, a practical algorithm instantiatingthe framework.
DiAMetR trains a population of meta-policies to be robust to different degrees ofdistribution shifts and then adaptively chooses a meta-policy to deploy based on the inferred test-timedistribution shift.
Our experiments verify the utility of adaptive distributional robustness undertest-time task distribution shift in a number of simulated robotics domains.
Related Work
Meta-reinforcement learning algorithms aim to leverage a distribution of training tasks to “learn areinforcement learning algorithm", that is able to learn as quickly on new tasks drawn from the samedistribution.
A variety of algorithms have been proposed for meta-RL, including memory-based[6, 22], gradient-based [9, 32, 11] and latent-variable based [30, 43, 42] schemes.
These algorithmsshow the ability to generalize to new tasks drawn from the same distribution, and have been appliedto problems ranging from robotics [24, 42, 15] to computer science education [39].
This line ofwork has been extended to operate in scenarios without requiring any pre-specified task distribution[10, 13] or in offline settings [5, 25, 23] making them more broadly applicable to a wider class ofproblems.
However, most meta-RL algorithms assume source and target tasks are drawn from thesame distribution, an assumption rarely met in practice.
Our work shows how the machinery ofmeta-RL can be made compatible with distribution shift at test time, using ideas from distributionalrobustness.
Some recent work shows that model based meta-reinforcement learning can be made tobe robust to a particular level distribution shift [20, 17] by learning a shared dynamics model againstadversarially chosen task distributions.
We show that we can build model-free meta-reinforcementlearning algorithms, which are not just robust to a particular level of distribution shift, but can adaptto various levels of shift.
Distributional robustness methods have been studied in the context of building supervised learningsystems that are robust to the test distribution being different than the training one.
The key ideais to train a model to not just minimize empirical risk, but instead learn a model that has thelowest worst-case empirical risk among an “uncertainty-set" of distributions that are boundedly closeto the empirical training distribution [33, 18, 2, 12].
If the uncertainty set and optimization arechosen carefully, these methods have been shown to obtain models that are robust to small amounts of distribution shift at test time [33, 18, 2, 12], finding applications in problems like federatedlearning [12] and image classification [18].
This has been extended to the min-max robustnesssetting for specific algorithms like model-agnostic meta-learning [3], but are critically dependent oncorrect specification of the appropriate uncertainty set and applicable primarily in supervised learningsettings.
Alternatively, several RL techniques aim to directly tackle the robustness problem, aimingto learn policies robust to adversarial perturbations [37, 41, 29, 28].
[40] conditions the policy onuncertainty sets to make it robust to different perturbation sets.
While these methods are able tolearn conservative, robust policies, they are unable to adapt to new tasks as DiAMetR does in themeta-reinforcement learning setting.
In our work, rather than choosing a single uncertainty set, welearn many meta-policies for widening uncertainty sets, thereby accounting for different levels oftest-time distribution shift.
Preliminaries
Meta-Reinforcement Learning aims to learn a fast reinforcement learning algorithm or a “metapolicy" that can quickly maximize performance on tasks T from some distribution p ( T ) .
Formally,each task T is a Markov decision process (MDP) M =
( S , A , P , R , γ, µ 0 ) ; the goal is to exploitregularities in the structure of rewards and environment dynamics across tasks in p ( T ) to acquireeffective exploration and adaptation mechanisms that enable learning on new tasks much faster thanlearning the task naively from scratch.
A meta-policy (or fast learning algorithm)
π meta maps a historyof environment experience h ∈ ( S ×
A × R ) ∗ in a new task to an action a , and is trained to acquireoptimal behaviors on tasks from p ( T ) within k episodes:
[Equation]
r ( i ) t (cid:35) , where h ( i ) t = ( s ( i ) 1: t , r ( i ) 1: t , a ( i ) 1: t − 1 ) ∪ ( s ( j ) 1: T , r ( j ) 1: T , a ( j ) 1: T )
i − 1 j =1 .
Intuitively, the meta-policy has two components: an exploration mechanism that ensures that appropriate reward signal is found for all tasks in the training distribution, and an adaptation mechanismthat uses the collected exploratory data to generate optimal actions for the current task.
In practice,the meta-policy may be represented explicitly as an exploration policy conjoined with a policyupdate[9, 30], or implicitly as a black-box RNN [6, 43].
We use the terminology “meta-policies"interchangeably with that of “fast-adaptation" algorithms, since our practical implementation buildson [27] (which represents the adaptation mechanism using a black-box RNN).
Our work focuseson the setting where there is potential drift between p train ( T ), the task distribution we have access toduring training, and p test ( T ) , the task distribution of interest during evaluation.
Distributional robustness [33] learns models that do not minimize empirical risk against the trainingdistribution, but instead prepare for distribution shift by optimizing the worst-case empirical riskagainst a set of data distributions close to the training distribution (called an uncertainty set ):
[Equation]
This optimization finds the model parameters θ that minimizes worst case risk l over distributionsq ϕ
( x ) in an ϵ -ball (measured by an f -divergence) from the training distribution p train ( x ) .
Distributionally Adaptive Meta-Reinforcement Learning
In this section, we develop a framework for learning meta-policies, that given access to a trainingdistribution of tasks p train ( T ) , is still able to adapt to tasks from a test-time distribution p test ( T ) thatis similar but not identical to the training distribution.
We introduce a framework for distributionallyadaptive meta-RL below and instantiate it as a practical method in Section 5.
4.1 Known Level of Test-Time Distribution Shift
We begin by studying a simplified problem where we can exactly quantify the degree to whichthe test distribution deviates from the training distribution.
Suppose we know that p test satisfies
[Table] [Table]
+JlAZKjQJfdwYUB2rWG4v
+deysM4QQjZdFEvERZG1jgUq8slMBQjTSiTXN9qsQGVlKGOrqhDcGZfnieNStk5KVduTkvVizyOAtknh+SYOOSMVMkVqZE6YeSRPJNX8mY8GS/Gu/ExbV0w8pk98gfG5w+965p/</latexit>
D ( p test ( T ) ||
p train ( T ))
<
ϵ for some ϵ > 0 , where D ( ·∥· ) is a probability divergence on the set of taskdistributions (e.g. an f -divergence [31] or a Wasserstein distance [36]).
A natural learning objectiveto learn a meta-policy under this assumption is to minimize the worst-case test-time regret across anytest task distribution q ( T ) that is within some ϵ divergence of the train distribution:
[Equation]
Solving this optimization problem results in a meta-policy that has been trained to adapt to tasksfrom a wider task distribution than the original training distribution.
It is worthwhile distinguishingthis robust meta-objective, which incentivizes a robust adaptation mechanism to a wider set of tasks,from robust objectives in standard RL, which produce base policies robust to a wider set of dynamicsconditions.
The objective in Eq 3 incentivizes an agent to explore and adapt more broadly, not actmore conservatively as standard robust RL methods [29] would encourage.
Naturally, the quality ofthe robust meta-policy depends on the size of the uncertainty set.
If ϵ is large, or the geometry of thedivergence poorly reflect natural task variations, then the robust policy will have to adapt to an overlylarge set of tasks, potentially degrading the speed of adaptation.
4.2 Handling Arbitrary Levels of Distribution Shift
In practice, it is not known how the test distribution p test deviates from the training distribution, andconsequently it is challenging to determine what ϵ to use in the meta-robustness objective.
We proposeto overcome this via an adaptive strategy: to train meta-policies for varying degrees of distributionshift, and at test-time, inferring which distribution shift is most appropriate through experience.
We train a population of meta-policies { π ( i ) meta } Mi =1 , each solving the distributionally robust meta-RLobjective (eq 3) for a different level of robustness ϵ i :
[Equation]
where ϵ M > ϵ M − 1 > . . .
>
ϵ 1 =(4)
In choosing a spectrum of ϵ i , we learn a set of meta-policies that have been trained on increasinglylarge set of tasks: at one end ( i = 1 ), the meta-policy is trained only on the original trainingdistribution, and at the other ( i = M ), the meta-policy trained to adapt to any possible task within theparametric family of tasks.
These policies span a tradeoff between being robust to a wider set of taskdistributions with larger ϵ (allowing for larger distribution shifts), and being able to adapt quickly toany given task with smaller ϵ (allowing for better per-task regret minimization).
With a set of meta-policies in hand, we must now decide how to leverage test-time experience todiscover the right one to use for the actual test distribution p test .
We recognize that the problemof policy selection can be treated as a stochastic multi-armed bandit problem (precise formulationin Appendix A), where pulling arm i corresponds to running the meta-policy π ϵ i meta for an entiremeta-episode ( k task episodes).
If a zero-regret bandit algorithm (eg: Thompson’s sampling [38]) isused , then after a certain number of test-time meta episodes, we can guarantee that the meta-policyselection mechanism will converge to the meta-policy that best balances the tradeoff between adaptingquickly while still being able to adapt to all the tasks from p test ( T ) .
To summarize our framework for distributionally adaptive meta-RL, we train a population of metapolicies at varying levels of robustness on a distributionally robust objective that forces the learned [Figure] adaptation mechanism to also be robust to tasks not in the training task distribution.
At test-time, weuse a bandit algorithm to select the meta-policy whose adaptation mechanism has the best tradeoffbetween robustness and speed of adaptation specifically on the test task distribution.
Combiningdistributional robustness with test-time adaptation allows the adaptation mechanism to work evenif distribution shift is present, while obviating the decreased performance that usually accompaniesoverly conservative, distributionally robust solutions.
4.3 Analysis
To provide some intuition on the properties of this algorithm, we formally analyze adaptive distributional robustness in a simplified meta RL problem involving tasks T g corresponding to reaching someunknown goal g in a deterministic MDP M , exactly at the final timestep of an episode.
We assumethat all goals are reachable, and use the family of meta-policies that use a stochastic exploratorypolicy π until the goal is discovered and return to the discovered goal in all future episodes.
Theperformance of a meta-policy on a task T g under this model can be expressed in terms of the statedistribution of the exploratory policy: Regret( π meta , T g )
=
1 d Tπ ( g ) .
This particular framework hasbeen studied in [10, 16], and is a simple, interpretable framework for analysis.
We seek to understand performance under distribution shift when the original training task distributionis relatively concentrated on a subset of possible tasks.
We choose the training distribution p train ( T g ) =(1 − β ) Uniform ( S 0 )
+ β Uniform ( S\S 0 ) , so that p train is concentrated on tasks involving a subsetof the state space {+S 
+}
⊂ S , with β a parameter dictating the level of concentration, and considertest distributions that perturb under the TV metric.
Our main result compares the performance of ameta-policy trained to an ϵ 2 -level of robustness when the true test distribution deviates by ϵ 1 . Proposition 4.1 Let ϵ i
=
min
{ ϵ i + β, 1 − |S
| |S| } .
There exists q ( T ) satisfying D TV ( p train , q ) ≤ ϵ 1where an ϵ 2 -robust meta policy incurs excess regret over the optimal ϵ 1 -robust meta-policy:
E q ( T )
[Regret( π ϵ 1 meta , T ) − Regret( π ϵ 2 meta , T )] ≥ (cid:18) c ( ϵ 1 , ϵ
[Equation]
The scale of regret depends on c ( ϵ 1 , ϵ 2 )
= (cid:113) ϵ 2 − 1 − 1 ϵ 1 − 1 − 1 , a measure of the mismatch between ϵ 1 and ϵ 2 .
We first compare robust and non-robust solutions by analyzing the bound when ϵ 2 = 0 .
In the regimeof β ≪ 1 , excess regret scales as O ( ϵ 1 (cid:113) 1 β ) , meaning that the robust solution is most necessarywhen the training distribution is highly concentrated in a subset of the task space.
At one extreme, ifthe training distribution contains no examples of tasks outside S 0 ( β = 0 ), the non-robust solutionincurs infinite excess regret ; at the other extreme, if the training distribution is uniform on the set ofall possible tasks ( β = 1 − |S 0 | |S| ), robustness provides no benefit .
We next quantify the effect of mis-specifying the level of robustness in the meta-robustness objective,and what benefits adaptive distributional robustness can confer.
For small β and fixed ϵ 1 , the excessregret of an ϵ 2 -robust policy scales as O ( (cid:113) max { ϵ 2 ϵ 1 , ϵ 1 ϵ 2 } ) , meaning that excess regret gets incurred ifthe meta-policy is trained either to be too robust ( ϵ 2 ≫ ϵ 1 ) or not robust enough ϵ 1 ≫ ϵ 2 .
Comparedto a fixed robustness level, our strategy of training meta-policies for a sequence of robustness levels π
meta ,θ , D Replay-Buffer ← do do 7:
Meta policy
Update π meta ,θ using off policy
2 [27] θ :=
θ + α ∇ θ
E z ∼ q ϕ
( z ) ( E π ϵ meta ,θ , P (1 k k (cid:88) i =
T (cid:88) t = r ω ( s ( i ) t , a ( i ) t , z ))) 8:
Adversarial task distribution: Update q ϕ using Reinforce [34] ϕ :
= ϕ − α ∇ ϕ
( E z ∼ q ϕ
( z ) [ E π ϵ meta ,θ , P [1 k k (cid:88) i =
T (cid:88) t = r ω ( s ( i ) t , a ( i ) t , z )]] + λD KL ( p train ( z ) ∥
q ϕ ( z )) 9:
Lagrange constraint multiplier: Update λ to enforce D KL ( p train ( z ) ∥ q ϕ ( z )) <
ϵ , λ := λ ≥ 0 λ
+ α ( D KL ( p train ( z ) ∥
q ϕ ( z ))
− ϵ ) 10: end for 11: end for { ϵ i }
Mi =1 ensures that this misspecification constant is at most the relative spacing between robustnesslevels: max
i ϵ i ϵ
i − 1 . This enables the distributionally adaptive approach to control the amount of excessregret by making the sequence more fine-grained, while a fixed choice of robustness incurs largerregret (as we verify empirically in our experiments as well).
DiAMetR: A Practical Algorithm for Meta-Distribution Shift
In order to instantiate our distributionally adaptive framework into a practical algorithm, we mustaddress how task distributions should be parameterized and optimized over, and also how the robustmeta-RL problem can be solved with stochastic gradient methods.
For simplicity, in the remainderof the paper, we focus on the setting where tasks share transition dynamics, but have differentreward functions.
We first introduce the individual components of task parameterization and robustoptimization, and describe the overall algorithm in Algorithm 1 and 2.
Parameterizing Task Distributions: Since we assume that variations in tasks correspond to changesin the reward function, the problem of representing a task distribution reduces to learning distributionsover reward functions.
We propose to learn a probabilistic model of the task reward functions seenin the training task distribution, and use the learned latent representation as a space on which toparameterize uncertainty sets over new task distributions.
Specifically, we jointly train a rewardencoder q ψ ( z | h ) that encodes reward samples from an environment history into the latent space, anda decoder r ω ( s, a, z ) mapping a latent vector z to a reward function using a dataset of trajectoriescollected from the training tasks.
This generative model over reward functions can be trained as astandard latent variable model by maximizing a standard evidence lower bound (ELBO), trading offreward prediction and matching a prior p train ( z ) (chosen to be the unit gaussian).
[Equation]
Having learned a latent space, we can parameterize new task distributions q ( T ) as distributions q ϕ ( z )(the original training distribution corresponds to p train ( z )
=
N (0 , I ) , and measure the divergencebetween task distributions as well using the KL divergence in this latent space D ( p train ( z ) ∥
q ϕ ( z )) .
Learning Robust Meta-Policies: Given this task parameterization, the next question becomes how toactually perform the robust optimization laid out in Eq:3.
The distributional meta-robustness objectivecan be modelled as an adversarial game between a meta-policy π ϵ meta and a task proposal distributionq ( T ) .
As described above, this task proposal distribution is parameterized as a distribution over latent [Table]
[Table] [Table] [Table] space q ϕ ( z ) , while π ϵ meta is parameterized a typical recurrent neural network policy as in [27].
Weparameterize { π ϵ i meta } Mi =1 as a discrete set of meta-policies, with one for each chosen value of ϵ .
This leads to a simple alternating optimization scheme (see Algorithm 1), where the meta-policy istrained using a standard meta-RL algorithm (we use off-policy RL 2 [27] as a base learner), and thetask proposal distribution with an constrained optimization method (we use dual gradient descent[26]).
Each iteration n , three updates are performed: 1) the meta-policy π meta updated to improveperformance on the current task distribution, 2) the task distribution q ( z ) updated to increase weight on tasks where the current meta-policy adapts poorly and decreases weight on tasks that the currentmeta-policy can learn, while staying close to the original training distribution, and 3) a penaltycoefficient λ is updated to ensure that q ( z ) satsifies the divergence constraint.
Algorithm 2 DiAMetR :
Meta-test phase 1: Given: p test ( T ) , Π = { π ϵ i meta ,θ } Mi = 2: Initialize TS = Thompson-Sampler() 3: for meta-episode n = 1 , 2 , ... do 4:
Choose meta-policy i =
TS.sample() 5:
Run π ϵ i meta ,θ for meta-episode 6:
TS.update(arm=
i, reward=meta-episode return) 7: end for
Test-time meta-policy selection: Since test-timemeta-policy selection can be framed as a multi-armedbandit problem, we use a generic Thompson’s sampling [38] algorithm (see Algorithm 2).
Each metaepisode n , we sample a meta-policy π ϵ meta with probability proportional to its estimated average episodicreward, run the sampled meta-policy π ϵ meta for anmeta-episode ( k environment episodes) and then update the estimate of the average episodic reward.
Since Thompson’s sampling is a zero-regret banditalgorithm, it will converge to the meta-policy thatachieves the highest average episodic reward andlowest regret on the test task distribution.
Experimental Evaluation
[Figure] We aim to comprehensively evaluate DiAMetR and answer the following questions: (1) Do metapolicies learned via DiAMetR allow for quick adaptation under different distribution shifts in thetest-time task distribution?
(2) Does learning for multiple levels of robustness actually help thealgorithm adapt more effectively than a particular chosen uncertainty level?
(3) Does proposinguncertainty sets via generative modeling provide useful distributions of tasks for robustness?
Setup.
We train DiAMetR on four continuous control environments: Wheeled navigation [11](Wheeled driving a differential drive robot), Ant navigation (Ant controlling a four legged roboticquadruped), Fetch reach and Block push [11] (Figures 4a to 4d) (see Appendix H for moredetails).
Each environment has a train task distribution T i ∼ p
train ( T ) such that each task T iparameterizes a reward function r
i ( s, a ) :
= r ( s, a, T i ) .
T i itself remains unobserved, the agent simply has access to reward values and executing actions in the environment.
The meta-policiesare evaluated on train task distribution p train ( T ) and on different distributionally shifted test taskdistribution { p i test ( T ) }
Ki =1 .
We use 4 random seeds for all our experiments and include the standarderror bars in our plots.
In all of these problems, the distribution of train and test tasks is determinedby the distribution of the underlying target locations s t , which determines the reward function (exactdistributions in Table 1).
Since these environments have sparse rewards, DiAMetR uses a structured
VAE to model reward distributions (see Appendix C for more details).
6.1 Adaptation to Varying Levels of Distribution Shift
During meta test, given a test task distribution p test ( T ) , DiAMetR uses Thompson sampling to selectthe appropriate meta-policy π ϵ meta ,θ within N = 250 meta episodes.
π ϵ meta ,θ can then solve any task T ∼ p test ( T ) within 1 meta episode ( k = 2 environment episode).
Since DiAMetR adaptivelychooses a meta-policy during test time, we compare it to RL 2 with test time finetuning.
Figure 5shows that RL 2 ’s performance more or less remains the same after test time finetuning showingthat 10 iteration (with 25 meta-episodes per iteration) isn’t enough for RL 2 to learn an meta-policyfor a new task distribution.
For comparison, RL 2 takes 1500 iterations (with 25 meta-episodes periteration) during training to learn a meta-policy for train task distribution.
To test DiAMetR’s ability to adapt to varying levels of distribution shift, we evaluate it on theabove mentioned test task distributions.
We compare DiAMetR with meta RL algorithms such as(off-policy) RL 2 [27], VariBAD [43] and HyperX [44].
Figure 6 shows that DiAMetR outperforms RL 2 , VariBAD and HyperX on test task distributions.
Furthermore, the performance gap between DiAMetR and other baselines increase as distribution shift between test task distribution and train taskdistribution increases.
Naturally, the performance of DiAMetR also deteriorates as the distributionshift is increased, but as shown in Fig 6, it does so much more slowly than other algorithms.
Wealso evaluate DiAMetR on train task distribution to see if it incurs any performance loss.
Figure 6shows that DiAMetR either matches or outperforms RL 2 , VariBAD, and HyperX on the train taskdistribution.
We refer readers to Appendix D for results on point-navigation environment and Appendix E for ablation studies and further experimental evaluations.
6.2 Analysis of Tasks Proposed by Latent Conditional Uncertainty Sets
[Figure]
We visualize the imagined test reward distribution for various distribution shifts.
Specifically,we create a heatmap of imagined test rewardfunctions.
Figure 7 visualizes the imagined testreward distribution in Ant-navigation environment in increasing order of distribution shiftswith respect to train reward distribution (withdistribution shift parameter
ϵ increasing fromleft to right).
The train distribution of rewardshas uniformly distributed target locations withinthe red circle.
As clearly seen in Figure 7, aswe increase the distribution shifts, the learnedreward distribution model imagines more targetlocations outside the red circle.
6.3 Analysis of Importance of Multiple Uncertainty Sets
DiAMetR meta-learns a family of adaptation policies, each conditioned on different uncertainty set.
As discussed in section 4, selecting a policy conditioned on a large uncertainty set would lead tooverly conservative behavior.
Furthermore, selecting a policy conditioned on a small uncertainty setwould result in failure if the test time distribution shift is high.
Therefore, we need to adaptivelyselect an uncertainty set during test time.
To validate this phenomenon empirically, we performedan ablation study in Figure 8.
As clearly visible, adaptively choosing an uncertainty set during testtime allows for better test time distribution adaptation when compared to selecting an uncertaintyset beforehand or selecting a large uncertainty set.
These results suggest that a combination oftraining robust meta-learners and constructing various uncertainty sets allows for effective test-timeadaptation under distribution shift.
DiAMetR is able to avoid both overly conservative behavior andunder-exploration at test-time.
Discussion
In this work, we discussed the challenge of distribution shift in meta-reinforcement learning andshowed how we can build meta-reinforcement learning algorithms that are robust to varying levelsof distribution shift.
We show how we can build distributionally “adaptive" reinforcement learningalgorithms that can adapt to varying levels of distribution shift, retaining a tradeoff between fastlearning and maintaining asymptotic performance.
We then show we can instantiate this algorithmpractically by parameterizing uncertainty sets with a learned generative model.
We empiricallyshowed that this allows for learning meta-learners robust to changes in task distribution.
There are several avenues for future work we are keen on exploring, for instance extending adaptivedistributional robustness to more complex meta RL tasks, including those with differing transitiondynamics.
Another interesting direction would be to develop a more formal theory providing adaptiverobustness guarantees in meta-RL problems under these inherent distribution shifts.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
Our work
is done in simulation and won’t have any negative societal impact.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes] This work does not actually use human subjects, and is done in simulation.
We have reviewed ethics guidelines and ensured that our paper conforms to them.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A] Math is used as a theory/formalism, but we don’t make any provable claims about it.
(b) Did you include complete proofs of all theoretical results?
[N/A
If you ran experiments...
(a) Did you include the code, data and instructions
needed to reproduce the main experimental results (either
in the
supplemental material or as a URL) [Yes
We have included the code along with a README in the supplemental material
(b) Did
you
specify all the training details (e g., data splits, hyperparameters, how they were chosen)?
[Yes] See Appendix I
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
[Table]
ments multiple times)?
[Yes] All plots were created with 4 random seeds with std error [Table]
bars.
[Table]
(d) Did you include the total amount of compute and the type of resources used (e.g., type [Table] of GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix I
[Table]
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
[Table]
(a) If your work uses existing assets, did you cite the creators?
[Yes] Environments weused are cited in section 6.
Codebase used are cited in Appendix I
[Table]
(b) Did you mention the license of the assets?
[N/A]
[Table]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
[Table]
We published the code and included all environments and assets as a part of this
[Table] (d) Did you discuss whether and how consent was obtained from people whose data you’re [Table]
using/curating?
[Yes] Environments and codebases we used are open-source.
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount [Table] spent on participant compensation?
[N/A]
[Table]