Distributionally Adaptive Meta Reinforcement Learning
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 shift of 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-time distribution shifts in the space of tasks.
Our framework centers on an adaptive approach to distributional robustness that trains a population of meta-policies to be robust to varying levels of distribution shift.
When evaluated on a potentially shifted test-time distribution of tasks, this allows us to choose the meta-policy with the most appropriate level of robustness, and use it to perform fast adaptation.
We formally show how our framework allows for improved regret under distribution shift, and empirically show its efficacy on simulated robotics problems under a wide range of distribution shifts.
Introduction
The diversity and dynamism of the real world require reinforcement learning (RL) agents that can quickly adapt and learn new behaviors when placed in novel situations.
Meta reinforcement learning 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 [40, 11, 34, 49].
Unfortunately, meta-RL agents assume tasks to be always drawn from the training task distribution and often behave erratically when asked to adapt to tasks beyond the training distribution [5, 8].
As an example of this negative transfer, consider using meta-learning to teach a robot to navigate to goals quickly (illustrated in Figure 1).
The resulting meta-policy learns to quickly adapt and walk to any target location specified in the training distribution but explores poorly and fails to adapt to any location not in that distribution.
This is particularly problematic for the meta-learning setting, since the scenarios where we need the ability to learn quickly are usually exactly those where the agent experiences distribution shift.
This type of meta-distribution shift afflicts a number of real-world problems including autonomous vehicle driving [10], in-hand manipulation [18, 1, 9], and quadruped locomotion [25, 23, 19], where training task distribution may not encompass all real-world scenarios.
In this work, we study algorithms that learn meta-policies resilient to task distribution shift at test time.
We assume the test-time distribution shift to be unknown but within a fixed range.
One approach to enable this resiliency is leveraging the distributional robustness framework [38] to learn a meta-policy robust to a wide range of distribution shifts.
However, the resulting meta-policy can be slow to adapt during test time.
In contrast, learning a meta-policy robust to a small range of distribution shifts enable faster test-time adaptation but leaves the meta-policy brittle to task distribution shifts.
Can we do better if we knew the degree of distribution shift apriori?
Yes, but since the test-time distribution shift is unknown, we would first want our algorithm to infer the degree of distribution shift and then deploy the appropriate meta-policy robust to the inferred degree of distribution shift.
To enable our approach, we use the distributional robustness framework to train meta-policies that prepare for distribution shifts by optimizing the worst-case empirical risk against a set of task distributions which lie within a bounded distance from the original training task distribution (known as an uncertainty set ).
This allows meta-policies to deal with potential test-time task distribution shift, bounding their worst-case test-time regret for distributional shifts within the chosen uncertainty set.
[Table]
t
[Table]
r a i
[Table]
t
[Table]
/HlZdKu15yzWv32vNq4KuIokSNyTE6JQy5Ig9yQJmkRRh7IM3klb8aT8WK8Gx/
z1hWjmDkkf2B8/gBBmJSl</latexit>
[Figure]
On meta-training tasks, π meta explores effec- tively and quickly learns the optimal behavior (top row).
When test tasks come from a slightly larger task distribution, exploration fails catastrophically, resulting in poor adap
Our key insight is that we can prepare for a variety of potential test-time distribution shifts by constructing and training 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 the most appropriate level of distributional robustness for the test distribution at hand.
We introduce a conceptual framework called distributionally adaptive meta reinforcement learning, formalizing this idea.
At train time, the agent learns robust meta-policies with widening uncertainty sets, preemptively accounting for different levels of test-time distribution shift that may be encountered.
At test time, the agent infers the level of distribution shift it is faced with, and then uses the corresponding meta-policy to adapt to the new task (Figure 2).
In doing so, the agent can adaptively choose the best level of robustness for the test-time task distribution, preserving the fast adaptation benefits of meta RL, while also ensuring good asymptotic performance under distribution shift.
We instantiate a practical algorithm in this framework called DiAMetR.
Our main contributions are twofold.
First, we show how to leverage the distributional robustness framework to make meta-reinforcement learning robust to a given level of distribution shift.
Secondly, we propose a framework for making meta-reinforcement learning resilient to a variety of task distribution shifts, and DiAMetR, a practical algorithm instantiating the framework.
Our experiments verify the utility of adaptive distributional robustness under test-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 a reinforcement learning algorithm", that is able to learn as quickly on new tasks drawn from the same distribution.
A variety of algorithms have been proposed for meta-RL, including memorybased [7, 26], gradient-based [11, 36, 14] and latent-variable based [34, 49, 48, 12] schemes.
These algorithms show the ability to generalize to new tasks drawn from the same distribution, and have been applied to problems ranging from robotics [28, 48, 19] to computer science education [44].
This line of work has been extended to operate in scenarios without requiring any pre-specified task distribution [13, 17], in offline settings [6, 29, 27] or in hard (meta-)exploration settings [50, 47], making them more broadly applicable to a wider class of problems.
However, most meta-RL algorithms assume source and target tasks are drawn from the same distribution, an assumption rarely met in practice.
Our work shows how the machinery of meta-RL can be made compatible with distribution shift at test time, using ideas from distributional robustness.
Some recent work shows that model based meta-reinforcement learning can be made to be robust to a particular level distribution shift [24, 21] by learning a shared dynamics model against adversarially chosen task distributions.
We show that we can build model-free meta-reinforcement learning algorithms, which are not just robust to a particular level of distribution shift, but can adapt to various levels of shift.
Distributional robustness methods have been studied in the context of building supervised learning systems that are robust to the test distribution being different than the training one.
The key idea is to train a model to not just minimize empirical risk, but instead learn a model that has the
[Table]
/zWCiura+sbxc3S1vbO7p6+f9ARUcIxaeOIRbzngSCMhqQtqWSkF3MCgcdI1xvfZH73nnBBo7AlJzFxAhiG1KcYpJJc
[Figure]
cgOQI4wsLQ1de0oIEOoiCpUH85cvWzWzBmMZWLlpIxyNF39yx5EOAlIKDEDIfqWGUsnBS4pZmRashNBYsBjGJK+oiEERDjp7IWpcaqUgeFHXFUojZn6eyKFQIhJ4KnO7GCx6GXif14/kf6Vk9IwTiQJ8XyRnzBDRkaWhzGgnGDJJooA5lTdauARcMBSpVZSIViLLy+TznnNuqjV7+rlxnUeRxEdoxNUQRa6RA10i5qojTB6RM
[Figure]
oFb1pT9qL9q59zFsLWj5ziP5A
[Figure]
lowest worst-case empirical risk among an “uncertainty-set" of distributions that are boundedly close to the empirical training distribution [38, 22, 2, 16].
If the uncertainty set and optimization are chosen carefully, these methods have been shown to obtain models that are robust to small amounts of distribution shift at test time [38, 22, 2, 16], finding applications in problems like federated learning [16] and image classification [22].
This has been extended to the min-max robustness setting for specific algorithms like model-agnostic meta-learning [3], but are critically dependent on correct specification of the appropriate uncertainty set and applicable primarily in supervised learning settings.
Alternatively, several RL techniques aim to directly tackle the robustness problem, aiming to learn policies robust to adversarial perturbations [42, 46, 33, 32].
[45] conditions the policy on uncertainty sets to make it robust to different perturbation sets.
While these methods are able to learn conservative, robust policies, they are unable to adapt to new tasks as DiAMetR does in the meta-reinforcement learning setting.
In our work, rather than choosing a single uncertainty set, we learn many meta-policies for widening uncertainty sets, thereby accounting for different levels of test-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 exploit regularities in the structure of rewards and environment dynamics across tasks in p (
T ) to acquire effective exploration and adaptation mechanisms that enable learning on new tasks much faster than learning the task naively from scratch.
A meta-policy (or fast learning algorithm)
π meta maps a history of environment experience h ∈ ( S ×
A × R ) ∗ in a new task to an action a , and is trained to acquire optimal behaviors on tasks from p (
T ) within k episodes:
[Equation] [Equation]
where h ( i ) t
= ( s ( i ) 1: t , r ( i ) 1: t , a
[Equation]
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 mechanism that 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 policy update[11, 34], or implicitly as a black-box RNN [7, 49].
We use the terminology “meta-policies" interchangeably with that of “fast-adaptation" algorithms, since our practical implementation builds on [31] (which represents the adaptation mechanism using a black-box RNN).
Our work focuses on the setting where there is potential drift between p train ( T ), the task distribution we have access to during training, and p test ( T ) , the task distribution of interest during evaluation.
Distributional robustness [38] learns models that do not minimize empirical risk against the training distribution, but instead prepare for distribution shift by optimizing the worst-case empirical risk
[Table]
[Figure]
k k
X i
[Figure]
T
X t
[Figure]
1 r (
[Figure] [Figure]
against 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 distributions q ϕ ( 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 training distribution of tasks p train ( T ) , is still able to adapt to tasks from a test-time distribution p test ( T ) that is similar but not identical to the training distribution.
We introduce a framework for distributionally adaptive 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 which the test distribution deviates from the training distribution.
Suppose we know that p test satisfies
D ( p test ( T ) ||
p train ( T ))
<
ϵ for some ϵ > 0 , where D ( ·∥· ) is a probability divergence on the set of task distributions (e.g. an f -divergence [35] or a Wasserstein distance [41]).
A natural learning objective to learn a meta-policy under this assumption is to minimize the worst-case test-time regret across any test 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 tasks from a wider task distribution than the original training distribution.
It is worthwhile distinguishing this 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 dynamics conditions.
The objective in Eq 3 incentivizes an agent to explore and adapt more broadly, not act more conservatively as standard robust RL methods [33] would encourage.
Naturally, the quality of the robust meta-policy depends on the size of the uncertainty set.
If ϵ is large, or the geometry of the divergence poorly reflect natural task variations, then the robust policy will have to adapt to an overly large 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, and consequently it is challenging to determine what ϵ to use in the meta-robustness objective.
We propose to overcome this via an adaptive strategy: to train meta-policies for varying degrees of distribution shift, 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-RL objective (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 increasingly large set of tasks: at one end ( i = 1 ), the meta-policy is trained only on the original training distribution, and at the other ( i = M ), the meta-policy trained to adapt to any possible task within the parametric family of tasks.
These policies span a tradeoff between being robust to a wider set of task distributions with larger ϵ (allowing for larger distribution shifts), and being able to adapt quickly to any 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 to discover the right one to use for the actual test distribution p test .
We recognize that the problem of policy selection can be treated as a stochastic multi-armed bandit problem (precise formulation in Appendix C), where pulling arm i corresponds to running the meta-policy π ϵ
i meta for an entire meta-episode ( k task episodes).
If a zero-regret bandit algorithm (eg: Thompson’s sampling [43]) is used , then after a certain number of test-time meta episodes, we can guarantee that the meta-policy selection mechanism will converge to the meta-policy that best balances the tradeoff between adapting quickly 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 adaptation mechanism to also be robust to tasks not in the training task distribution.
At test-time, we use a bandit algorithm to select the meta-policy whose adaptation mechanism has the best tradeoff between robustness and speed of adaptation specifically on the test task distribution.
Combining distributional robustness with test-time adaptation allows the adaptation mechanism to work even if distribution shift is present, while obviating the decreased performance that usually accompanies overly 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 some unknown goal g in a deterministic MDP M , exactly at the final timestep of an episode.
We assume that all goals are reachable, and use the family of meta-policies that use a stochastic exploratory policy π until the goal is discovered and return to the discovered goal in all future episodes.
The performance of a meta-policy on a task T g under this model can be expressed in terms of the state distribution of the exploratory policy: Regret( π meta , T g )
=
1 d Tπ ( g ) .
This particular framework has been studied in [13, 20], and is a simple, interpretable framework for analysis.
We seek to understand performance under distribution shift when the original training task distribution is 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 subset of the state space
S 0 ⊂ S , with β a parameter dictating the level of concentration, and consider test distributions that perturb under the TV metric.
Our main result compares the performance of a meta-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 ) ≤ ϵwhere an ϵ 2 -robust meta policy incurs excess regret over the optimal ϵ 1 -robust meta-policy:
[Equation]
The scale of regret depends on c ( ϵ 1 ,
[Equation]
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 regime of β ≪ 1 , excess regret scales as O ( ϵ(cid:113) 1 β ) , meaning that the robust solution is most necessary
] θ := θ + α ∇ θ E T ∼ q ϕ ( T )
[ E π ϵ meta ,θ , P T [1 k k (cid:88) i =
T (cid:88) t = r T ( s ( i ) t , a ( i ) t )]] 8:
Adversarial distribution: q ϕ using Reinforce [39] ϕ :=
ϕ − α ∇ ϕ
( E T ∼ q ϕ ( T )
[ E π ϵ meta ,θ , P T [1 k k (cid:88) i =
T (cid:88) t = r T ( s ( i ) t , a ( i ) t )]] + λD KL ( p train ( T ) ∥
q ϕ ( T ))) 9:
Lagrange constraint multiplier: Update λ enforce D KL ( p train ∥ q ϕ ) <
ϵ , λ := λ ≥ 0 λ
+ α ( D KL ( p train ( T ) ∥
q ϕ ( T ))
− ϵ ) 10: for when the training distribution is highly concentrated in a subset of the task space.
At one extreme, if the training distribution contains no examples of tasks outside S 0 ( β = 0 ), the non-robust solution incurs infinite excess regret ; at the other extreme, if the training distribution is uniform on the set of all 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 excess regret of an ϵ 2 -robust policy scales as O ( (cid:113) max { ϵ 2 ϵ 1 , ϵ 1 ϵ 2 } ) , meaning that excess regret gets incurred if the meta-policy is trained either to be too robust ( ϵ 2 ≫ ϵ 1 ) or not robust enough ϵ 1 ≫ ϵ 2 .
Compared to a fixed robustness level, our strategy of training meta-policies for a sequence of robustness levels { ϵ i } Mi =1 ensures that this misspecification constant is at most the relative spacing between robustness levels:
max i ϵ i
ϵ
i − 1 .
This enables the distributionally adaptive approach to control the amount of excess regret by making the sequence more fine-grained, while a fixed choice of robustness incurs larger regret (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 must address how task distributions should be parameterized and optimized over .
We must also address how the robust meta-RL problem can be solved with stochastic gradient methods.
We first introduce the individual components of task parameterization and robust optimization, describe the overall algorithm in Algorithm 1 and 2, and visualize components of DiAMetR in Fig 3.
5.1 Parameterizing Task Distributions
Handling in-support distribution shifts:
For handling in-support task distribution shifts, we propose to represent new task distributions as re-weighted training task distribution q ( T ) ∝ w T p train ( T ) where w
T > 0 is a parameter.
Since we have a finite set of training tasks, say {T i } n tr i =1 , new task distributions become q ( T i )
=
w T i (cid:80) n tr i =1 w T i .
With a slight abuse of notation, we can write empirical training task distribution as p train ( T i ) =
1 n tr .
We can use KL divergence to measure the divergence between training task distribution and test task distribution D ( p train ( T ) ∥
q ϕ ( T ))
= (cid:80) n tr i =1 1 n tr log
[Equation]
We collectively represent the parameters ϕ
= { w T i } n tr i =1 .
Using this parameterization, the training objective (equation 3) becomes
[Equation]
Handling out-of-support distribution shift: For handling out-of-support task distribution shifts, we propose to learn a probabilistic model of the training task distribution, and use the learned latent representation as a space on which to parameterize uncertainty sets over new task distributions.
Specifically, we jointly train a task encoder q ψ ( z | h ) that encodes an environment history into the latent space, and a decoder T ω ( s, a, z ) mapping a latent vector z to a property of the task using a dataset of trajectories collected from the training tasks.
To describe the exact form of T ω , we consider how tasks can differ and list two scenarios: (1) Tasks differ in reward functions: T ω takes form of reward functions r ω ( s, a, z ) that maps a latent vector z to a reward function and (2) Tasks differ in dynamics: T ω takes form of dynamics p ω ( s, a, z ) that maps a latent vector z to a dynamics function.
This generative model can be trained as a standard latent variable model by maximizing a standard evidence lower bound (ELBO), trading off reward prediction and matching a prior p train ( z ) (chosen to be the unit gaussian).
[Equation]
Having learned a latent space
[Equation]
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 divergence between task distributions as well using the KL divergence in this latent space D ( p train ( z ) ∥
q ϕ ( z )) .
Using this parameterization, the training objective (equation 3) becomes
[Equation]
when rewards differ
[Equation]
when dynamics differ
[Equation]
5.2 Training and test-time selection of meta-policies
Learning Robust Meta-Policies: Given this task parameterization, the next question becomes how to actually perform the robust optimization laid out in Eq:3.
The distributional meta-robustness objective can be modelled as an adversarial game between a meta-policy π ϵ meta and a task proposal distribution q ( T ) .
As described above, this task proposal distribution is parameterized as a distribution over latent space q ϕ ( z ) , while π ϵ meta is parameterized a typical recurrent neural network policy as in [31].
We parameterize { π ϵ 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 is trained using a standard meta-RL algorithm (we use off-policy RL 2 [31] as a base learner), and the task proposal distribution with an constrained optimization method (we use dual gradient descent [30]).
Each iteration n , three updates are performed: 1) the meta-policy π meta ,θ updated to improve performance on the current task distribution, 2) the task distribution q ϕ ( T ) updated to increase weight on tasks where the current meta-policy adapts poorly and decreases weight on tasks that the current meta-policy can learn, while staying close to the original training distribution, and 3) a penalty coefficient λ is updated to ensure that q ϕ ( T ) satisfies the divergence constraint.
[Figure]
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 meta ,
TS.
(arm=
i, reward=meta-episode return) 7: end for
Test-time meta-policy selection: Since test-time meta-policy selection can be framed as a multiarmed bandit problem, we use a generic Thompson’s sampling [43] algorithm (see Algorithm 2).
Each meta-episode n , we sample a meta-policy π ϵ meta with probability proportional to its estimated average episodic reward, run the sampled meta-policy π ϵ meta for an meta-episode ( k environment episodes) and then update the estimate of the average episodic reward.
Since Thompson’s sampling is a zero-regret bandit algorithm, it will converge to the meta-policy that achieves the highest average episodic reward and lowest regret on the test task distribution.
Experimental Evaluation
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 the test-time task distribution?
(2) Does learning for multiple levels of robustness actually help the algorithm adapt more effectively than a particular chosen uncertainty level?
(3) Does proposing uncertainty sets via generative modeling provide useful distributions of tasks for robustness?
Setup.
We use DiAMetR on four continuous control environments: Ant navigation (controlling a four-legged robotic quadruped), Wind navigation [6] (controlling a linear system robot in presence of wind), Object localization (controlling a Fetch robot to localize an object through its gripper) and Block push (controlling a robot arm to push an object) [14] (Figures 4a to 4d) (see
Appendix D for details about reward function and dynamics).
We design various meta RL tasks from these environments.
Each meta RL task has a train task distribution T i ∼ p
train ( T ) such that each task
T i either parameterizes a reward function r i ( s, a ) :
= r ( s, a, T i ) or a dynamics function p
i ( s, a ) :
= p ( s, a, T i ) .
T i itself remains unobserved, the agent simply has access to reward values and executing actions in the environment.
The learned meta-policies are evaluated on different distributionally shifted test task distributions { p i test ( T ) }
Ki =1 which are either in-support or out-of- support of training task distribution.
In all meta RL tasks, the train and test task distribution is determined by the distribution of an underlying task parameter (i.e. wind velocity w T for Wind navigation and target location s T for other environments), which either determines the reward function or the dynamics function.
While tasks vary in reward functions for Ant navigation ,
Object localization and Block push , they vary in dynamics function for Wind navigation (exact task distributions in Table 2).
We use 4 random seeds for all our experiments and include the standard error bars in our plots.
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 select the appropriate meta-policy π ϵ meta ,θ within N = 250 meta episodes.
π ϵ meta ,θ can then solve any new task T ∼ p test ( T ) within 1 meta episode ( k environment episode).
To test DiAMetR’s ability to adapt to varying levels of distribution shift, we evaluate it on different test task distributions, as detailed in Table 2.
We compare DiAMetR with meta RL algorithms such as (off-policy) RL 2 [31],
VariBAD [49] and HyperX [50].
Since DiAMetR uses 250 meta-episodes to adaptively choose a metapolicy during test time, we finetune RL 2 , VariBAD and HyperX with 250 meta-episodes of test task distribution to make the comparisons fair (see Appendix G for the finetuning curves).
Figure 5 show that DiAMetR outperforms RL 2 , VariBAD and HyperX on out-of-support and in-support shifted test task distributions.
Furthermore, the performance gap between DiAMetR and other baselines increase as distribution shift between test task distribution and train task distribution increases.
Naturally, the performance of DiAMetR also deteriorates as the distribution shift is increased, but as shown in
Fig 5, it does so much more slowly than other algorithms.
We also evaluate DiAMetR on train task distribution to see if it incurs any performance loss.
Figure 5 shows that DiAMetR either matches or outperforms RL 2 , VariBAD, and HyperX on the train task distribution.
6.2 Analysis of Tasks Proposed by Latent Conditional Uncertainty Sets
We visualize the imagined test reward distribution through their heatmaps for various distribution shifts.
To generate an imagined reward function, we sample z ∼ q ϕ ( z ) and pass the z into r ω ( s, a, z ) =
r ω ( s, z ) (given the learned reward is only dependent on state as mentioned in Appendix A.1).
We then take the ant agent and reset its ( x, y ) location to different points in the (discretized) grid [ − 1 , 1] 2 and calculate r ω ( s, z ) at all those points.
This gives us a reward map for a single imagined reward function.
We sample 10000 of these reward functions and plot them together.
Figure 6 visualizes the imagined test reward distribution in Ant-navigation environment in increasing order of distribution shifts with respect to train reward distribution (with distribution shift parameter ϵ increasing from left to right).
The train distribution of rewards has uniformly distributed target locations within the red circle.
As seen in Figure 6, the learned reward distribution model imagines more target locations outside the red circle as we increase the distribution shifts.
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 to overly conservative behavior.
Furthermore, selecting a policy conditioned on a small uncertainty set would result in failure if the test time distribution shift is high.
Therefore, we need to adaptively
[Figure] [Figure]
select an uncertainty set during test time.
To validate this phenomenon empirically, we performed an ablation study in Figure 7.
As clearly visible, adaptively choosing an uncertainty set during test time allows for better test time distribution adaptation when compared to selecting an uncertainty set beforehand or selecting a large uncertainty set.
These results suggest that a combination of training robust meta-learners and constructing various uncertainty sets allows for effective test-time adaptation under distribution shift.
DiAMetR is able to avoid both overly conservative behavior and under-exploration at test-time.
Discussion
In this work, we introduce a distributionally “adaptive" meta RL framework for tackling task distribution shifts and argue for adaptation in face of these distribution shifts.
There are several avenues for future work we are keen on exploring, for instance extending adaptive distributional robustness to more complex meta RL tasks, including those with image observations.
Another interesting direction would be to develop a more formal theory providing adaptive robustness guarantees in meta-RL problems under these inherent distribution shifts.
Acknowledgements
The authors thank the members of Improbable AI Lab & RAIL for discussions and helpful feedback.
We thank MIT Supercloud and the Lincoln Laboratory Supercomputing Center for providing compute resources.
This research was supported by an NSF graduate fellowship, a DARPA Machine Common
Sense grant, a MURI grant from the Army Research Office under the Cooperative Agreement
Number W911NF-21-1-0097, and an MIT-IBM grant.
This research was also partly sponsored by the United States Air Force Research Laboratory and the United States Air Force Artificial
Intelligence Accelerator and was accomplished under Cooperative Agreement Number FA8750-19-21000.
The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the United States
Air Force, the United States Army Research Office or the U.S. Government.
The U.S. Government is authorized to reproduce and distribute reprints for Government purposes, notwithstanding any copyright notation herein.
[-This leads to learning of two weighting factors λ 
-]
[-{ ϵ i + β, 1 − |S 
-]
[-q ′ 
-]