Distributionally Adaptive Meta Reinforcement Learning
Authors: Anurag Ajay
Abhishek Gupta
Dibya Ghosh
Sergey Levine
Pulkit Agrawal
Venue: NeurIPS 2022 Accept (2022)
Keywords: Meta reinforcement learning
Distributional robustness

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.

Distributionally Adaptive Meta Reinforcement Learning
Anurag Ajay" fe$, Abhishek Gupta *t8, Dibya Ghosh", Sergey Levine", Pulkit Agrawal
Improbable Al Lab"
MIT-IBM Watson Al Lab²
University of California, Berkeley
Massachusetts Institute Technologys
Abstract
Meta-reinforcemmnt learning algorithms provide a data-driven way to acquire poli- cies that quickly adapt to many tasks with varying rewards dynamics functions. However, learned meta-policies are often effective only on the exact task distribu- tion which they were trained and struggle in the presence of distribution shift test-time rewards or transition dynamics. In this work, we develop a frame- work for meta-RL algorithms that are able to behave appropriately under test-time distribution shifts the space tasks. Our framework centers on an adaptive approach to distributional robustness that trains population of meta- vetaa-pllicies be robust to varying levels of distribution shift. When evaluated on a potentially shifted test-time distribution of tasks, this allows to choose the meta-policy with the most appropriate level of robustness, and use it to perform fast adaptation. We formally show how 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 framework for conferring this ability to RL agents, by learning a "meta- "meta-policy" trained to adapt as quickly as possible to tasks from a provided training distribution [40, 11, 34, 49]. Unfortunately, meta-RL 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 example of this negative transfer, consider using meta-learning to teach robot 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 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, 9], and quadruped locomotion [25, 23, 19), where training task distribution may not encompass all real-world scenarios.
In this work, study algorithms that learn meta- meta-policae policies resilient task distribution shift 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 meta-policy robust to a wide range of distribution shifts. However, the resulting meta-policy can be slow to adapt
denotes equal contribution. Correspondence to aajayemit. edu and abhguptatcs. washington. edu
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

during test time. In contrast, learning a meta-policy robust to 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 set of task distributions which lie within bounded distance from the original training task distribution (known as an uncertainty set). This allows meta- peet--poicies to deal with potential rest-time task distribution shift, bounding their worst-case test-time regret for distributional shifts within the chosen uncertainty set.
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.
Episode
Episode
Meta Train
Tmeta
By preparing for adaptation against each of these uncer- tainty sets, an agent is able to adapt to variety of poten- tial test-time distribution shifts by adaptively choosing the most appropriate level of distributional robustness for the test distribution at hand. We introduce a conceptual frame- work called distributionally adaptive meta reinforcement learning, formalizing this idea. At train time, the agent learns robust meta- veta-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 adap- tively 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 practical algorithm in this framework called DiAMetR.
Meta Test
1
Tmeta
Figure 1: Failure of Typical Meta-RL. On meta-training tasks, Tmea explores effec- tively and quickly learns the optimal behav- ior (top row). When test tasks come from a slightly larger task distribution, exploration fails catastrophically. resulting in poor adap- tation behavior (bottom row).
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, practical algorithm instantiating the framework. Our experiments verify the utility of adaptive distributional robustness under test-time task distribution shift in number of simulated robotics domains.
2 Related Work
Meta-reinforcemeent learning algorithms aim to leverage a distribution of training tasks to "learn reinforcement learning algorithm". that able to learn as quickly on new tasks drawn from the same distribution. A variety of algorithms have been proposed for meta-RL, RL, including memory- based [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 jexporaixoo settings [50, 47], making them more broadly applicable to a wider class of problems. However, most meta-RI. 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 RL can be made compatible with distribution shift at test time, using ideas from distributional robustness. Some recent work shows that model based meta-reinffrremment 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- -ree meta-reinforcemeet 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 model to not just minimize empirical risk, but instead learn model that has the

---

Tulo, Castbuston
6
Imaggned
Buffer
Trai bakks
Meta-train on train-task distribution
Meta-train Imagined test-task
Meta-policy selection during meta- test
Meta-policy adaptation
distributions
During meta-train, DiAMetR learns meta-policy T and task distribution model Tu (s. 2) on train task distribution. Then, it uses the task distribution model to imagine different shifted test task distributions on which it learns different meta potaapooiiiess /im each corresponding to different level of robustness. During meta-test, it chooses an appropriate meta-policy based on inferred test task distribution shift with Thompson's 's sampling and then quickly adapts the selected meta-policy to individual tasks.
Figure 2:
2-22
lowest worst-case empirical risk among an "uncertainty-sst 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, 32j. [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 unable to adapt meta-reinforcemmen 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.
tasks as DiAMetR does in the
3 Preliminaries
Learning aims learn fast reinforcement learning algorithm "meta- policy" that can quickly maximize performance on tasks T from some distribution p(7). Formally, each task T is Markov decision process (MDP) M (S, A, P, R, 7, Po); 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) meeta maps history of environment experience h E (S A x R)' new task to an action a, and is trained to acquire optimal behaviors on tasks from p(T) within & episodes:
min ET TTT) (Regret (Tmtta: T)), Tima
T
Regret(Tmetan T) [(T]] - E
J( = max ER,T
Σ re]
( Mucu 18 (1) )
K
where h = (5 La o
(1) 1:7
( ( 1-11
1:7h 1:7/=11
Intuitively, the meta-policy has two components: exploration mechanism that ensures that appro- priate reward signal is found for all tasks 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 policy update[11. 1, 34], implicitly interchangeably with that of "fast-adaptation algorithms, since our practical implementation builds on [31] (which represents the adaptation mechanism using blx RNN). Our work focuses on the setting where there is potential drift between Prrain (T), 7, the task distribution we have access to during training, and Prest (T). ) the task distribution of interest during evaluation.
black-box RNN [7, 49]. We use the terminology "meta-policies"
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
3

---

Meta Policy
h
MDP
St+t1
F1
Meta Policy Selection
F4
π meta.
O2
T2-1
Constrained
Fuccto,
Amets.
task distribution D(Prrain (T (T)) 6
Phenl)
Meta Train phase
Meta Test phase
Figure 3: During meta-train phase, DiAMetR learns a family of meta-policies robust to varying levels of distribution shift (as characterized by e.). During meta-test phase, given potentially shifted test-time distribution of tasks, DiAMetR chooses the meta-policy with the most appropriate level of robustness and use perform fast adaptation for new tasks sampled from the same shifted test task distribution.
against a set of data distributions close to the training distribution (called an uncertainty set):
min max 1 (z) (x) 0)] s.t. Diprain (x)main())leer (z)ll S
This optimization finds the model parameters that minimizes worst case risk / over distributions (z) in e-ball (measured by an f-divergence) from the training distribution Preain (z).
4 Distributionally Adaptive Meta-Reinforcement Learning
In this section, we develop a framework for learning meta- peta-eolicies, that given access to training distribution tasks Pusin (T), still able adapt tasks from test-time distribution Pust that is similar but not identical to the training distribution. We introduce framework for distributionally adaptive meta-RI. below and instantiate 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 Phest satisfies D(pres (T) Praiin (T)) < for some >> 0, where D(-Il) is a probability divergence on the set of task distributions (e.g. f-divergence [35] or Wasserstein distance [41]). A natural learning objective to learn meta-policy under this assumption is to minimize the worst-case test-time regret across any test task distribution (T) that within some divergence of the train distribution:
min R (Tmmsaa Puain (T), e),
Tasta
R( (mota Purain (T),)) = max E ETg)()) [Regret(Tmen (meta T) TTT)
s.t. Dp (T) (l))) Se
(3)
Solving this optimization problem results in 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 robust adaptation mechanism from robust objectives in standard RL, which produce base policies robust to a wider set of dynamics conditions. The objective in Eq 3 incentivizes agent to explore and adapt more broadly, not act more conservatively as standard robust RL methods
wider set of tasks,
[33] would encourage. Naturally, the quality of the robust meta-policy depends the size of the uncertainty set. If is large, the geometry of the divergence poorly reflect natural task variations, then the robust policy will have to adapt to an overly large of tasks, potentially degrading the speed of adaptation.
4.2 Handling Arbitrary Levels of Distribution Shift
In practice, not known how the test distribution Phes 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 { meen j== objective (eq 3) for a different level of robustness Ex:
J. each solving the distributionally robust meta-RL.
arg min R metta Pirain (T), [2)
where EM EM- ...
(4)
mmeta
i=1
In choosing spectrum of é;, we learn set of meta-policies that have been trained on increasingly large set tasks: at one end (i = 1), the meta-policy trained only on the original training distribution, and at the other (i = .M), the meta-policy trained to adapt 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 c (allowing for better per-task regret minimization).
With set of meta- peta--ooiiies in hand, we must now decide how to leverage test-time experience to discover the right one to use for the actual test distribution Plest- We recognize that the problem of policy selection can be treated as stochastic multi- bandit problem (precise formulation in Appendix C), where pulling arm corresponds to running the meta-policy T for an entire meta-episode (k: task episodes). If zero-regret bandit algorithm (eg: Thompson's sampling [43]) is used, then after a certain number of test-time meta episodes, wc can guarantee that the meta- meta-policy selection mechanism will converge to the meta-policy that best balances the tradcoff between adapting quickly while still being able to adapt to all the tasks from Prest (
To summarize our framework for distributionally adaptive meta-RL, we train a population of meta- policies 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 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 distribu- tional robustness in simplified meta RL problem involving tasks T, corresponding to reaching some unknown goal 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, under this model can be expressed in terms of the state distribution of the exploratory policy: Regret (Tmem; T₂) = (.) This particular framework has been studied in [13, 20], and is 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 Purain (T,) = (1 - 3)nnfform(( (S,) Unnformm (S', (S),,,, of the state space S. C S, with 3 parameter dictating the level of concentration, and consider test distributions that perturb under the TV metric. Our main result compares the performance of meta-policy trained to an level of robustness when the true test distribution deviates by 61.
that Purin is concentrated
tasks involving subset
Proposition 4.1. Let minfq + 3, [So]]. There exists (T)) satisfying Drv (Poain: 9) S c where an 62 ea-rotust meta policy incurs excess regret over the optimal 61 e1-robust meta meta-policy:
E (Rerott (nina T) Regret(m t T) 2
(2(c1, + d(e,, 62) -2
(5)
(1(116))) -S )
(6)
The seale regret depends on c((
measure of the mismatch between 6 and cg.
We first compare robust and non-robust solutions by analyzing the bound when = 0. In the regime of B < 1, excess regret scales as O(er/7)) meaning that the robust solution is most necessary
5

---

Algorithm DiiAMee::Meta-rriinnng phase
1: Given: Prain (T), Return: maaa )
2: "meta.0. Butter Solve equation with off off-policy RL
3: Prior Punin(7) t Solve eq 8 using Buspla
4: for in (e,,..... ) do
5: Initialize qe- " and À for iteration n = 1, 2.. do Meta-policy: Update Timeaa using off-policy RI. [31]
7:
k T
( (
rT
1=1 t=1
8:
Adversarial task distribution: Update do using Reinforce [39]
T
a = aV (V men, PT
rT o |] + XDxi. (Prain (aalaannnTT44)) (T)lla
9:
Lagrange constraint multiplier: Update to enforce Dxi (pusin (10) < c,
À ==2 X + DxD (Prain (T) -c)
10: end for
11:
end 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 So (3 = 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 (B3
[So]] 3
robustness provides 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 3 and fixed 61, the excess regret of e2-robust policy scales as Og1 max
52 fi }), meaning that excess regret gets incurred if Cq <g
the meta-policy is trained either be robust (e2 e1) or robust enough e > 62. Compared to fixed robustness level, our strategy of training meta-policies for a sequence of robustness levels { ensures that this misspecification constant is most the relative spacing between robustness levels: max; _t. This enables the distributionally adaptive approach to control the amount of excess regret by making the sequence more fine-grained, while 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. 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: to represent new task distributions re-weighted training task distribution (T)) Prrain (T) where WT > 0 is a parameter. Since we have a finite set of training tasks, say {T new task distributions WIT, With a slight abuse of notation, we can write empirical training task
For handling in-support task distribution shifts, we propose
become ((T,) 2e Whr," (Tr)
distribution as Purin
1. We can use KL divergence to measure the divergence between Eis wT,
training task distribution and test task distribution ((Prain())
Ei Fhy netti
log
We collectively represent the parameters {WT, 1h Using this parameterization, the training

---

objective (equation 3) becomes
d
T
noT Wi::
max min E Puiia)) Braan P.P
Σ
1=
DKL (Phain (()(((()) S
Handling out-of-support distribution shift: we propose to learn a probabilistic model of the training task distribution, and use the learned latent representation space on which parameterize uncertainty
For handling out-of-support task distribution shifts,
over new task distributions. Specifically, we jointly train a task encoder Qy (=|h) that encodes an environment history into the latent space, and a decoder (s, a, mapping a latent vector
property of the task using dataset of trajectories collected from the training tasks. To describe the exact form of Tw. we consider how tasks can differ and list reward functions (s, a, z) that maps latent vector 2 to reward function and (2) Tasks differ in dynamics: Tu takes form of dynamics P= (s, a, 2) that maps latent vector to a dynamics function. This generative model can be trained as standard latent variable model by maximizing a standard evidence lower bound (ELBO), trading off reward prediction and matching a prior Prrain (z) (chosen to be the unit gaussian).
scenarios: (I) Tasks differ in reward functions: Tw takes form of
Σ 1(T at, 2), (t), 2), (xL m)
min EAND w.
E (elg
((Tu(se; ae h, (rw(se, av re)² when rewards differ I(T,(11, 2), h = (8,, (1, s111 when dynamics differ
Having learned a latent space, we can parameterize new task distributions (T) as distributions qo (2) (the original training distribution corresponds to Punin (O, I). and measure the divergence between task distributions as well using the KI. divergence in this latent space (ptrain (lg Using this parameterization, the training objective (equation 3) becomes
T
I]
max min Eg (x)
E
when rewards differ
8
k
1=1
t=1
max min E. (2)())
E oramee
when dynamics differ
(Puman (Prainnii))) (z))(g(()) Se
(9)
5.2 Training and test-time selection of 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 an adversarial game between meta-policy =fnets and task proposal distribution q(T). As described above, this task proposal distribution is parameterized as distribution over latent space (2), while T is parameterized typical recurrent neural network policy in [31]. We parameterize {m 1M j== as discrete set of meta-policies, with one for each chosen value of
Learning Robust
This leads simple alternating optimization scheme (see Algorithm 1), where the meta-policy is trained using a standard meta-RL. algorithm (we use off- off-policy RL2 [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 meta-policy A updated to improve performance on the current task distribution, 2) the task distribution (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 the original training distribution, and 3) a penalty coefficient À is updated to ensure that (T) satisfies the divergence constraint.
7

---

C
(a) Ant navigation
(b) Wind navigation
(c) Object localization
(d) Block push
Figure 4: The agent needs to either navigate in absence of winds, navigate in presence of winds, use its gripper to unobserved target location or push the block
localize object by green sphere cube), by exploring its environment functions for Ant navigation, Object localizati on and Block push, they vary dynamics function for Wind navigation.
unobserved target location, indicated experiencing reward. While tasks vary reward
Algorithm 2 DiAMetR: Meta-test phase
1: Given: Pesu 11 = {m 1
2: Initialize TS Thompson-Sani -Sampeer()))
3: for meta-episode n 1, ... do
4: Choose meta-policy i TS. sample ()
5: Run T for meta-episode
6: TS. update(
arm=i,
reward=meta-episode return)
7: end for
Test-time meta-policy selection: Since test-time meta-policy selection be framed multi- armed bandit problem, we use a generic Thompson's 's sampling [43] algorithm (see Algorithm 2). Each meta-episode n, we sample a meta-policy meees with probability proportional to its estimated average episodic reward, run the sampled meta-policy Teets for an meta-episode (k environment episodes) and then update the estimate of the average episodic reward. Since Thompson's '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 comprehensively evaluate DiAMetR and answer the following questions: (1) Do meta- policies learned via DiAMetR allow for quick adaptation under different distribution shifts in the test-time task distribution? algorithm adapt more effectively than a particular chosen uncertainty level? (3) Does proposing uncertainty sets generative modeling provide useful distributions of tasks for robustness?
(2) Does learning for multiple levels of robustness actually help the
Setup. We use DiAMetR on four continuous control environments: Ant navigation (controlling four-legged robotic quadruped), Wind navigation [6] (controlling 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 object) [14] (Figures 4a to 4d) (see Appendix for details about reward function and dynamics). We design various meta RL tasks from these environments. Each meta RI task has train task distribution T; Purain (Tr such that each task T, either parameterizes a reward function (x, a) = r(s, a, TT) or a dynamics function pi(s, a) == p(s, a, T.). Tr 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 (T)) which are either in-support 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 wy for Wind navigation and target location sy for other environments), which either determines the reward function or the dynamies function. While tasks vary 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 random seeds for all our experiments and include the standard error bars in our plots.
8

---

Wind
navigation
Object Localization
Black
Sucrzo Ray
Suzrzo Raze
tepex
310
1.00
037
612
0.5
0.15
112
Speal
Raillo
Radio
novigation
Objeect Localization
Black Push
X13
Sapcos Rate
Sapcos Rate
Ricleet iider Hepux
300
3.50
0.57
0.55
0.55
0.50
0.55
0.55
0.
00
Spead
Radio
Figure 5: We evaluate DiAMetR and meta RL algorithms (RL3, VariBAD and HyperX) different in-support (top row) and sut-of-support (bottom row) shifted test task distributions. DiAMeIR either matches or outperforms RL", VariBAD and HyperX on these task distributions. The first point pı on the horizontal axis indicates the task parameter (A) distribution (4(0, p1) and the subsequent points p. indicate task parameter (A) distribution U(p.-1,P.). While task parameter for wind navigation is wind velocity wy, is target location sy for other environments. Table details the task distributions used in this evaluation.
6.1 Adaptation to Varying Levels of Distribution Shift
During meta test, given a test task distribution Pies (T))(T), DiAMetR uses Thompson sampling to select the appropriate meta-pot veta--poiccy TH m within N = 250 meta episodes. T mesaa can then solve any new task T ~ Prea(() within meta episode (A 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 [31], VariBAD [49] and HyperX [50]. Since DiAMetR uses 250 meta-episodes to adaptively choose a meta- policy during test time, we finetune RL2 VariBAD and HyperX with 250 meta- 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", VariBAD and HyperX 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 RL2, VariBAD, and HyperX on the train task distribution.
ont-of--uppor and in-support shifted test
6.2 Analysis of Tasks Proposed by Latent Conditional Uncertainty Sets
We visualize the imagined test reward distribution through their heatmaps for various distribu- tion shifts. To generate an imagined reward function, we sample z (z) and pass the into (w, a, =) = (s, (s. (given the learned reward is only dependent on state as mentioned in Ap- pendix A. We then take the ant agent and reset its (z, y) location different points in the (discretized) grid 1] and calculate (s, 2) 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 natt-navigatin on 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 policy conditioned on small uncertainty set would result in failure if the test time distribution shift is high. Therefore, need to adaptively

---

( 0.1 (b) = 0.2 ( 0.4
(d) 0.8
Figure 6: Imagined test reward distributions in Ant-navigation environment increasing order distribution shifts. Train reward distribution is uniform within the red circle.
Wind
eavigation
Object Locelization
Black Push
pais
Sucves Raze
oo sova pea
Sucves Mare
EANAAD Jabe?
305
1.5
0.7 Wind
060
3.54
6:55
0.30
965
311
612
0.44
0.5
Radie
navigation
Object Localization
Block Pash
10.2
Succco Rai.
A.rk oka
XL1
KL
3nd
210
0.40
Sponl
Figure 7: Adantively choosing an uncertainty for DiAMetR policy (Adapt) during test time allows better adapt to test time distribution shift than choosing an uncertainty set beforehand (Mid). Choosing a large uncertainty set for DiAMetR policy (Conservative) leads to a conservative hehavior and hurts its performance when test time distribution shift is low. While top row contains in-support task distribution shifts, the bottom contains out-of-support task distribution shifts.
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 large uncertainty set. These results suggest that combination of training robust meta-learners and constructing various uncertainty sets allows for effective tist-time adaptation under distribution shift. DiAMetR is able to avoid both overly conservative behavior and under-exploration at test-time.
7 Discussion
In this work, we introduce a distributionally "adaptive" meta RL framework for tackling task distribu- tion 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 W91 1991N2221110000 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-2. 1000. 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 authorized to reproduce and distribute reprints for Government purposes, notwithstanding any copyright notation herein.
10

---

References
[1] T. Chen, J. Xu, and P. Agrawal. system for general in-hand object re-orientation. In A. Faust, D. Hsu, and G. Neumann, editors, Proceedings of the 5th Conference Robot Learning. volume 164 of Proceedings of Machine Learning Research, pages 297-307. PMLR, 08-11 Nov 2022. URL https: http://///oceceng..e..o mml. prwsp/ w111//11112222 html.
[2] J. Cohen, E. Rosenfeld, and Z. Kolter. Certified adversarial robustness via randomized smooth- ing. In International Conference on Machine Learning, pages 1310-1320. PMLR, 2019.
[3] L. Collins, A. Mokhtari, and Shakkottai. Distributioon-agoosic model-agnostic meta-learning. CoRR, abs/2002.04766, 2020. URL
04766.
[4] P.-T. De Boer, D. Kroese, S. Mannor, and R. Rubinstein. tutorial on the cross-entropy method. Annals of operations research, 134(1):19667, 2005.
[5] T. Deleu and Bengio. The effects of negative adaptation in model-agnostic meta-learning. arXiv preprint arXiv: /8/2.02/9/ 2018.
[6] R. Dorfman, Shenfeld, and A. Tamar. Offline meta learning of exploration. arXiv preprint arXiv:2008.02:3 2020.
[7] Y. Duan, J. Schulman, X. Chen, P. L. Bartlett, I. Sutskever, and P. Abbeel. RI2: Fast reinforce- ment learning via slow reinforcement learning. arXiv preprint arXiv:1611.0279) 2016.
[8] A. Fallah, A. Mokhtari, and A. Ozdaglar. Generalization of model-agnostic meta-learning algorithms: Recurring and unseen tasks. Advances in Neural Information Processing Systems, 34, 2021.
[9] N. Fazeli, A. Ajay, and A. Rodriguez. Long-horizen prediction and uncertainty propagation with residual point contact learners. In 2020 IEEE International Conference on Robotics and Automation (ICRA), pages 7898-7904, 2020. doi: 10. 1091019//949945...... 0B0H//CCAA4444422200999991111111 2020.9196511.
[10] A. Filos, P. Tigkas, R. McAllister, N. Rhinehart, S. Levine, and Y. Gal. Can autonomous vehicles identify, recover from, and adapt to distribution shifts? In International Conference on Machine Learning, pages 3145-3153. PMLR, 2020.
[11] C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference machine learning, pages 1126-1135. PMLR, 2017.
[12] H. Fu, H. Tang, J. Hao, C. Chen, Feng, D. Li, and Liu. Towards effective context for meta-reinforcement learning: an approach based on contrastive learning. In Procedingss of the AAAI Conference on Artificial Intelligence, volume 35, pages 7457-7465, 2021.
[13] A. Gupta, B. Eysenbach, C. Finn, and S. Levine. Unsupervised meta-learning for reinforcement learning. arXiv preprint arXiv: 04640, 2018.
[14] A. Gupta, R. Mendonca, Y. Liu, Abbeel, and S. Levine. Meta-reinforcement learning of structured exploration strategies. Advances in neural information processing systems, 31, 2018.
[15] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. arXiv preprint arXiv: arXiv:/801.0// /801.0/290, 2018.
[16] J. Hong. H. Wang, Z. Wang, and J. Zhou. Federated robustness propagation: Sharing adversarial robustness in federated learning. arXiv preprint arXiv:2106.10196, 2021.
[17] A. Jabri, K. Hsu, A. Gupta, B. Eysenbach, S. Levine, and C. Finn. Unsupervised curricula for visual meia-reeinorreement learning. Advances in Neural Information Processing Systems, 32, 2019.
[18] I. Ke, J. Wang. T. Bhattacharjee, B. Boots, and S. Srinivasa. Grasping with chopsticks: Combating covariate shift in model-free imitation learning for fine manipulation. In 202/ IEEE International Conference on Robotics and Automation (ICRA), pages 6185-6191. IEEE, 2021.
11

---

[19] A. Kumar, Z. Fu, D. Pathak, and J. Malik. Rma: Rapid motor adaptation for legged robots. arXiv preprint arXiv:2107.044 07.04444,44 2021.
[20] L. Lee, B. Eysenbach, Parisotto, E. Xing, Levine, and R. Salakhutdinov. Efficient exploration via state marginal matching. CoRR, abs/1906.05274, 2019. URL http://arxiv. org/abs/ haaab//996...... 11119966.. 05274.
[21] Z. Lin, G. Thomas, G. Yang, and T. Ma. Model-based adversarial meta-reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, ed- itors, Advances Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL
weciiinngnnneuuuree nww.inuurp.... httpsssccpppp///22//////// / Abstract html.
[22] A. Madry, A. Makelov, L. Schmidt, D. Tsipras, and A. Vladu, Towards deep learning models resistant to adversarial attacks. arXiv preprint arXiv: arXiv:1706. 1706.06066...33 2017.
[23] G. B. Margolis, G. Yang, K. Paigwar. T. Chen, and P. Agrawal. Rapid locomotion via reinforce- ment learning, arXiv preprint arXiv:2205.0258 2022.
[24] R. Mendonca, X. Geng, C. Finn, and S. Levine. Meta-reinforcement learning robust to distribu- tional shift via model identification and experience relabeling. CoRR, abs/2006.071 070006 2020. URL https: / /ttp/aaxvvv////////////e///aaav. org/ abs 200os/2200000000 07178.
[25] T. Miki, J. Lee, J. Hwangbo, L. Wellhausen, V. Koltun, and M. Hutter. Learning robust perceptive locomotion for quadrupedal robots in the wild. Science Robotics, 7(62).eabk2822, 2022.
[26] N. Mishra, M. Rohaninejad, X. Chen, and P. Abbeel. A simple neural attentive meta-learner. arXiv preprint arXiv: 2017.
[27] E. Mitchell, R. Rafailov, X. B. Peng, S. Levine, and C. Finn. Offline meta-reinforcemennt learning with advantage weighting. In International Conference on Machine Learning, pages 7780-7791. PMLR, 2021.
[28] A. Nagabandi, I. Clavera, S. Liu, R. S. Fearing, P. Abbeel, S. Levine, and C. Finn. Learning to adapt in dynamic, real-world environments through meta-reinforcemmnnt learning. arXiv preprint OrXiv::800s / /347, 2018.
[29] A. Nair, A. Gupta, M. Dalal, and S. Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006. 09359, 2020.
[30] Y. Nesterov. Primal-dual subgradient methods for convex problems. Mathematical programmingg 120(1):221-259, 2009.
[31] T. Ni, B. Eysenbach, S. Levine, and R. Salakhutdinov. Recurrent model-free RI. a strong base- line for many POMDPs, 2022. URL https:// nett/forrummddd500000000000
[32] T. P. Oikarinen, W. Zhang, A. Megretski, L. Daniel, and T. Weng. Robust deep reinforcement learning through adversarial loss. In M. Ranzato, A. Beygelzimer, Y. N. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6- 14. 2021, virtual, pages 26156-261617 2021. URL https: / http///pooeeeen....../ neurips.cccc Abstrrtbbst.a....ee.e..c. html.
[33] L. Pinto, J. Davidson, R. Sukthankar, and A. Gupta. Robust adversarial reinforcement learning. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, ICML. 2017, Sydney, NSW, Australia, 6-11 I August 2017, volume 70 of Proceedings of Machine Learning Research, pages 2817-2826. PMLR, 2017. URL http: www.eressss.pepm//eeeeemeemeeese
/ //proceedings.n mlr. prwssprssssssp.. Tivess/v70/piie p70/pinto17a.ll.
[34] K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen. Efficient off- off-policy meta-reinforcement learning via probabilistic context variables. In International conference on machine learning, pages 5331-5340. PMLR, 2019.
12

---

[35] A. Rényi. On measures of entropy and information. In Proceedings of the Fourth Berkeley Symposium on Mathematical Statistics and Probability, Volume Contributions the Theory of Statistics, volume 4, pages 547-562. University of California Press, 1961.
[36] Rothfuss, D. Lee, Clavera, T. Asfour, and P. Abbeel. Promp: Proximal meta- meta-policy search. arXiv preprint arXiv:/800.06700 2018.
[37] Schulman, Wolski, Dhariwal, A, Radford, and O. Klimov. Proximal policy optimiza- tion algorithms. CoRR, wws//100..0633. 2017. URL http://dblp... ami-trien de/dd///bb..
j jouralls/corr/cor1///000 htm Hth.:mmmsshhhmmaakWRR
[38] A. Sinha, H. Namkoong, R. Volpi, and J. Duchi. Certifying some distributional robustness with principled adversarial training. arXiv preprint arXiv: 10571, 2017.
[39] R. S. Sutton, D. McAllester, S. Singh, and Y. Mansour. Policy gradient methods for reinforce- ment learning with function approximation. Advances in neural information processing systems, 12, 1999.
[40] S. Thrun and L. Pratt, editors. Learning to Learn. Springer, 1998. ISBN 978- 1-4613-7527-2. doi: 01077997811-615-55222 URL https://doi. org/10. 1101000.1000100
[41] L. N. Vaserstein. Markov processes over denumerable products of spaces, describing large systems of automata. Problemy Peredachi Informatsii, 5(3):64-72 12.64,2,,11 1969.
[42] E. Vinitsky, Y. Du, K. Parvate, K. Jang, P. Abbeel, and A. M. Bayen. Robust reinforcement learning using adversarial populations. CoRR, abs/2008.01822. 2020. URL https://arxiv org/ hrg/abs/0008 01825.
[43] D. Wolpert and W. Macready. No free lunch theorems for optimization. IEEE Transactions on Evolutionary Computation, ((1):67-82, 1997. doi: 10. 10.1109/4235.585889
[44] M. Wu, N. Goodman, C. Piech, and C. Finn. Prototransformer: A meta-leaming approach to providing student feedback. arXiv preprint arXhv:2107.14033 2021.
[45] A. Xie, S. Sodhani, C. Finn, J. Pineau, and A. Zhang. Robust policy learning over multiple uncertainty sets. arXiv preprint arXiv:2202.07013. 2022.
[46] H. Zhang, H. Chen, D. S. Boning, and C. Hsieh. Robust reinforcement learning on state observations with learned optimal adversary. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL htt.httts///oonreiiiiie www.eeew..eeeeeeeeeeeeeeeeeeeeee http://orrmmddddc2/aaaaa
[47] J. Zhang, J. Wang, H. Hu, T. Chen, Y. Chen, C. Fan, and C. Zhang. Metacure: Meta reinforce- ment learning with empowerment-drrven exploration. In International Conference on Machine Learning, pages 12600-12610, PMLR, 2021.
[48] Z. Zhao, A. Nagabandi, K. Rakelly, C. Finn, and S. Levine. Meld: Meta-reinfooremment learning from images via latent state models. arXiv preprint 1XXiv:2010.13995 2020.
[49] L. Zintgraf, K. Shiarlis. M. Igl, Schulze, Y. Gal, K. Hofmann, and Whiteson. Varibad: A very good method for bayes-adaptive deep rl via meta-learning. arXiv preprint 188.0.9900088348. 2019.
[50] L. M. Zintgraf, L. Feng, C. Lu, M. Igl, K. Hartikainen, K. Hofmann, and S. Whiteson. Ex- ploration in approximate hyper-state space for meta reinforcement learning. In International Conference on Machine Learning, pages 12991-13001 PMLR, 2021.
13

---

Checklist
1. 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 7
(c) Did you discuss any potential negative societal impacts of your work? [N/A] Our work 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.
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A] Math used as theory/formallsmm but we don't It make any provable claims about it.
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main ex- perimental results (either in the supplemental material as 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
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes] All plots were created with 4 random seeds with std error bars.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes] See Appendix
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] Environments we used are cited in section Codebase used are cited in Appendix
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either the supplemental material or as URL? [Yes] We published the code and included all environments and assets as part of this
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [Yes] Environments and codebases we used are open-source.
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "yjXUvjUtQPY", "reviewer": "Reviewer_Asdp", "summary": "This paper proposes a new framework (called DiAMetR) for meta RL, aiming at improving the ability of meta policy in adapting to varying levels of task distribution shift. This is achieved by leveraging the distributional robust framework to re-formulate the meta RL objective to optimize the worst-case empirical risk under a uncertainty set of test-time task distribution. To be specific, the task distribution is parameterized through modeling the latent representation of the reduced reward functions. Then, the training alternates between 1) the training of meta policy against multiple imagined task distributions, each of which has a different distribution deviation constraint, and 2) the training of the imagined task distributions in an adversarial manner.\n\nThe performance of DiAMetR is evaluated in four goal-reaching environments, demonstrating the superiority in adaptation to varying levels of task distribution shift when compared with RL$^2$, VariBAD and HyperX. Ablation study and analysis are also provided.\n", "questions": "1) How do the authors think about the generalization ability of the reward function models trained with samples collected by changing policies under the original task distributional? (as my comment says in the part of Strengths and Weaknesses)\n2) For Line 7 in Algorithm 1, i.e., training the meta policy with imagined task distributions, how are the state-action trajectories collected? I can imagine three options: 1) directly interact with environment since the transition dynamics is the same; 2) interact with a learned transition dynamics model; 3) reuse the experiences in the replay buffer.\n3) How are the set of task distribution deviation (i.e., $\\{\\epsilon\\}$) chosen for each environments in the experiments? I did not found the associated content (I also checked the appendix), did I miss it?\n4) Do the baselines VariBAD and HyperX also use finetuning with the same number of samples during test-time adaptation?\n5) Although the authors provide Figure 8, it would be useful and informative to provide the analysis between the arm (i.e., some $\\epsilon$) adaptively chosen by Algorithm 2 and the oracle test distribution.\n\n&nbsp;\n\nMinors:\n- The legends in Figure 5,6,8 can be enlarged for better clarity.\n\n&nbsp;\n\n\nI may improve my rating if my concerns and questions are well addressed.\n", "limitations": "The major limitation is the proposed method focuses on reward variations while assuming transition dynamics unchanged. This is considered as one of the future directions in Section 7.\n\n&nbsp;\n\n\nI think there are three more limitations which are not well discussed in this paper:\n1) The first one is the choice of the uncertainty sets (i.e., the multiple $\\epsilon$). Intuitively, these hyperparameters are significant to the learning performance of the proposed method, and one may need several attempts or prior knowledge to determine the task-dependent choices.\n2) The second one is that the performance of current method in dense-reward environments, or non-goal reaching environments remains unclear. I also notice that the structured VAE is kind of specialized for sparse-reward tasks.\n3) The third one is my concern on the sample cost of test-time adaptation (i.e., 250 meta episode). I am a little reluctant to agree with the ‘quick adaptation’, especially when considering other possible environments with horizons longer than e.g. 50, 200 (Table 2 in the Appendix) used in this work.\n", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n+ The paper is well written and the presentation of the algorithm is very clear.\n+ The related work is almost sufficient.\n+ The main idea is general and novel in meta RL. Insightful formal analysis is provided.\n+ The proposed method is evaluated from multiple perspectives, from performance comparison to visual analysis.\n\n&nbsp;\n\nCons:\n\n- The proposed method is kind of limited.\n  - This is mainly because DiAMetR relies on the assumption that  the target tasks only differ at reward function while the transition dynamics remain unchanged. Personally, I think this cripples the significance of the proposed method a lot.\n  - In addition, it is unclear how DiAMetR will perform in dense-reward environments or non-goal reaching environments. It seems that the specific model structure (i.e., structured VAE in appendix is specialized to sparse-reward environments).\n- The latent representation and generative model of reward function (Eq.6) implicitly depends on the distribution of state-action pairs. In practice, the distribution of state-action pairs is generated by the policies optimized during the learning process under the original training task distribution.\n  - I have no question on the worst-case objective (Eq.3 and Eq.4) when oracle imagined task distribution and task generation are considered. However, since the latent representation and generative model of reward function are trained with samples collected by changing policies under the original task distributional, I worry about the ability of the generalization of such reward function models, in turn the effectiveness and rationality, especially in more possible environments.\n- Important experimental details are missing (see my questions below).\n\n&nbsp;\n\nThere are some related works which can be included for more complete empirical evaluation and related work survey:\n\n- MetaCURE: Meta Reinforcement Learning with Empowerment-Driven Exploration. ICML 2021\n- Towards Effective Context for Meta-Reinforcement Learning: an Approach based on Contrastive Learning. AAAI 2021\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "m14LCIY1Ds", "reviewer": "Reviewer_Z27V", "summary": "The paper proposes a meta-RL algorithm for dealing with test-time distribution shift.\nThis is achieved by learning a population of policies that have different levels of distributional robustness by an adversarial algorithm.\nAt test time, a policy is chosen amongst the population by an efficient bandit algorithm.\nThe paper demonstrates empirically that this way of adapting to test-time distribution achieves better performance than competing methods.\nTheoretical result provides a bound on the difference of expected regret of a mis-specified robustness epsilon.\n", "questions": "# Major Concerns\n\n* Equation 1 is expected regret, shouldn’t the regret be written as  - expected return of the meta-policy and the actual term in the expectation at the moment is just an MC estimate, this would seem more natural to me, because the outer expectation is taken over different tasks, not environment dynamics, so the expected return of the meta-policy doesn’t even appear in the minimization problem. \n\n* Fix fonts in figure 5, they are way too small. Also, RL2 and RL2 finetune legend doesn’t match the dashed lines. It’s not clear to me what is diameter oracle here? If it’s the best policy that can be chosen from the set of trained meta-policies, why is then the diameter adapt mean higher than the oracle’s mean? This figure needs better clarification.\n\n* It seems to me that your results hinge on the fact that if you are lucky enough that the test distribution is close to the train distribution, you should be less conservative, otherwise you should be more conservative and therefore sacrifice performance. Doesn’t it make more sense to have policies that are less conservative (train on task distributions of smaller support, slightly overlapping support) and then choose amongst them? In this way there would be no ‘conservativeness tradeoff’. This is not shown in figure 7.\n\n*  Ideally I would also like to see the distribution over policies from the population of meta-policies that get picked for solving tasks. Is it skewed towards more conservative policies? The plot should show that conservativeness grows with epsilon, currently you only have the result that picking the most conserative policy isn’t the best thing. (figure 8).\n\n* l.141 I don’t understand why equation 3 incentivizes exploration as you say it? I understand how it leads to distributional robustness, since you want to minimize the maximum regret, which is the difference to the optimal expected return. But how does this encourage exploration? I don’t see any reason why this should be the case…\n\n* How many adaptation policies are used in the experiments? Experiments need better clarification in general.\n\n# Minor Concerns\n\n* Section 6.2, by test reward distributions you mean q(z | h)? Please put inline math somewhere in this section to make the clear where the samples come from? What projection method was used to generate these heatmaps in 2d? Or is z actually 2-dimensional?\n\n* Figure 8: fix the fonts, same as figure 5. I don’t understand how is diameter mid uncertainty set chosen here exactly, so it’s hard for me to infer anything from this result, what does it imply?\n\n* Equation 2, why is q parametrized here? It should be just a set of distributions that are close in terms of f-divergence to the training distribution p, there is no reason why they should be parameterized.\n\nRequirements to increase my score:\n\n* Better explanation of the experiments section (what do the plots mean).\n* Better explanation of section 6.2 (projection method, where do the samples come from).\n* Argue clearly why does it make sense to have more and less conservative policies in comparison to less conservative ones with different task-space specialisations.\n* Answer all the listed points\n\n\n", "limitations": "The authors should make it clear how much it costs them to learn this population of policies and how it relates to the test-time distribution shift. I imagine that the number of policies is also going to affect the performance in the end, because it's directly related to the granularity of epsilon (which controls conservativeness).", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths\n\n* The method proposed in the paper is novel (as far as I know) and seems to show improvement in comparison to other meta RL methods.\n* The presentation of the method is clear, I didn't have many problems following the storyline\n\n# Weaknesses\n\n* Missing details in the experiment section.\n* Simplified task spaces where goals are be sampled from uniform distribution and limitation to only shift in reward function as result, not transition dynamics for example.\n* It is unclear to me based on the presentation how conservative the policies that are chosen by the bandit are in the end.* “the distributional meta-robustness objective can be modelled as an adversarial game between the policy and the task distribution”. Unfortunately this is the only sentence that hints at a game-theoretic treatment of the topic, but it doesn’t exist in the paper. I like that using an adversarial formulation for inferring the task sampling distribution q allows for a lot of flexibility regarding task distributions, the environments though used in the experiments to validate the method all deal with relatively simple distribution shifts concerning the reward (just uniform distributions of goals), so perhaps the tasks used for evaluation are too simple?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "HXIRSyFZ3z", "reviewer": "Reviewer_V7F9", "summary": "In meta RL, the most common assumption is that the training task distribution is the same as the test task distribution, which may not hold in practice. To address the unknown distribution shift in practice, this paper proposes an adaptive method, by learning a set of policies with varying levels of distribution shift. Under a certain level of distribution shift, the policy is learned by leveraging the framework of distributional robustness.  At test time, the selection of policies is formulated as a MAB problem. ", "questions": "1. I was wondering why the authors did not capture the distribution over reward and transition functions, in their method to parameter the task distribution. It seems not very difficult. \n2. Assuming p_train(z) = N(0, I) does not seem reasonable to me. N(0, I) is prior. Even training the ELBO  only on training task distribution, the posterior is usually different from the prior. \n3. Is it possible to approve some uni-modality property of this special MAB problem? If so, you can use a more efficient method with the unimodal bandit. ", "limitations": "1. The method can be extended to capture both variations in reward function and transition function. \n2. The theoretical results are limited to a toy environment. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n1. The distribution shift problem is very common in practice. \n2. The idea of learning a set of policies with varying levels of distribution shift, when the test distribution cannot be observed in advance, is quite reasonable. \n3. The paper is well-organized. \n\nWeaknesses:\n1. The theoretical analysis is only limited to a toy environment.\n2. The proposed method only parameters the task distribution by parameterizing distribution over the reward function. However, different tasks usually differ in both reward and transition functions.\n3. The test-time meta policies selection is based on MAB. I was wondering if it is possible to improve this MAB method, perhaps by considering some uni-modality properties of this special MAB problem(Is it possible to approve that?). ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "GdctoKfYKwS", "reviewer": "Reviewer_qqiy", "summary": "This paper identifies and addresses the problem of the distributional shift in test task distributions for meta-RL agents, which often struggle under such circumstances. For this matter, the work leverages the distributional robustness framework to alleviate this effect, by training a population of meta-policies in different levels of distribution shift. In order to do so, the proposed method (DiAMetR) firstly learns a single meta-policy in the train distribution and uses the collected trajectories to learn a probabilistic latent space to parameterize task distributions. Then, each subsequent meta-policy is learned in an adversarial game between itself and a task proposal distribution with some degree of shift. At test-time, a bandit algorithm based on Thompson sampling selects the appropriated meta-policy (level of distribution shift) to interact with the test task. The proposed method is evaluated in 4 different continuous control environments and presents better adaptation performance in the scenarios of test distribution shift when compared to strong baselines. The work also presents ablations to validate the components of the method and a theoretical proposition to better understand a meta-policy performance under distribution shift.", "questions": "The main suggestions or questions/concerns to be clarified are pointed out in the Strengths/Weaknesses and Limitations sections.", "limitations": "First of all, I appreciate that the work raised the limitation regarding handling only tasks with reward variation (as opposed to dynamics variation). It would be interesting to discuss ways to address this limitation in future work, in light of the proposed framework.\n\nBesides this limitation, there is also another point worth discussing:\n\nThe distributional shift phenomenon can happen in many different ways. In other words, considering the training distribution, there might have several different test distributions with the same level of shift (given a divergence metric). It is hard to know if a single “adversarial game” per shift level would be able to cover enough possibilities of test task distributions.\n\nFor the presented experimental setup, this does not seem to be a problem, as the shift would happen in only one dimension (reward) that varies across very few variables (e.g, radius and angle of the goal). However, in more complex cases, this can be concerning, and adding new adversarial games incurs new meta-policies, which also increases adaptation complexity (which is already a concern as pointed out in the weaknesses). Curious to know your thoughts on this point.\n\nFigure 7d illustrates this concern: it is possible to see that the imagined test distributions are more “skewed” in the west-south directions, but the real test distribution could be the opposite, for example.\n", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n- Addressing distributional shift for meta-RL agents is a very important topic, especially since OOD generalization is a property that the community is actively looking for understanding and improvements;\n- The paper is very organized and well written, all ideas are clear and the work is appropriately placed in prior literature;\n- The method is novel, technically very sound, and brings a solid contribution to the field of meta-RL towards the mitigation of distributional shift in test task distribution;\n- The theoretical framework is clearly stated and the derived proposition is relevant; the experimental setups support the main claims and back up the proposed method. As a strong positive point, there is a helpful, extensive ablation and the raised questions during the experimental setup are very relevant to validating the method.\n\nWeaknesses/Concerns:\n\n- Using a multi-armed bandit for meta-policy selection seems to be very costly (in terms of test-time sample complexity), and the complexity grows with the number of meta-policies. For instance, the example from Section 6.1 shows that  DiAMetR requires 250 meta-episodes (or 500 episodes) to complete adaptation. Prior literature shows evidence for few-episode adaptation, sometimes online adaptation [1, 2]. I understand that these methods are not focused on distributional shift, but in scenarios where test-time sample complexity is also a requirement, this performance seems concerning.\n- In the same line as the previous point, it would be interesting to show the results from variBAD and hyperX in the scenario of test time adaptation (Figure 5 and Appendix G). It would be helpful to understand how these other methods would handle this trade-off between performance and adaptation complexity for the distributional shift scenario.\n- As noted by the paper, the proposed implementation is limited by meta-RL environments comprising only variations in the reward distribution, missing those ones with variations in the dynamics. This is an important concern, as there is no way to validate if the proposed method is still valid for more complex benchmarks. For instance, it would be interesting to see how it performs in MetaWorld ML10/ML45 [3].\n\nDespite these weaknesses, the understanding is that the method and empirical validation present satisfactory quality for acceptance, given the proposed scope.\n\nMinor issues:\n\n- Figure 2 is not cited and actually does not seem to have any support in the text besides Algorithm 1.\n- A diagram with all the modules and training/test stages would be very helpful to improve clarity and presentation. For example, something like Figure 2 in the [VariBAD paper](https://arxiv.org/pdf/1910.08348.pdf).\n- In Equation 1, $J(\\pi_{t})$ is not defined (Is the minus sign a typo?)\n- Algorithm 1, Line 3 cites Equation 7, but this actually refers to the meta-policy selection, is this right?\n- In Appendix F, the summary of proof (L575): the paper describes a 3-stage proof, but stage “1)” is repeated twice.\n\nReferences:\n\n[1] Zintgraf et. al. VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning. ICLR, 2020.\n\n[2] LC Melo. Transformers are Meta-Reinforcement Learners. ICML, 2022.\n\n[3] Yu et. al. Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning. CoRL, 2019.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept