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 [38, 9, 32, 46].
Unfortunately, meta-RL agents are prone to overfitting to the distribution of tasks they are trained on and have been shown to behave erratically when asked to adapt to tasks beyond the training distribution [4, 7].
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.
Overfitting 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 [8], in-hand manipulation [15, 1], and quadruped locomotion [23, 21, 17], where the test-time task distribution may not be well represented during training.
In this work, we study meta-RL algorithms that learn meta-policies resilient to task distribution shift at test time.
One approach to enable this resiliency is to leverage the framework of distributional Decision Awareness in Reinforcement Learning Workshop at the 39th International Conference on Machine Learning ( robustness [36], training 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 (often referred to 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.
However, choosing an appropriate uncertainty set can be quite challenging without further information about the test environment, significantly impacting the test-time performance of algorithms under distribution shift.
Large uncertainty sets allow resiliency to a wider range of distribution shifts [Table] but the resulting meta-policy adapts very slowly at test time; smaller uncertainty sets enable faster test-time adaptation, but leave the meta-policy brittle to task distribution shifts.
Can we get the best of both worlds?
[Table]
/HlZdKu15yzWv32vNq4KuIokSNyTE6JQy5Ig9yQJmkRRh7IM3klb8aT8WK8Gx/
z1hWjmDkkf2B8/gBBmJSl</latexit> ⇡ meta <latexit sha1_base64="04ZwISNQzfwUIERNw0RDchdyb+M=">AAAB+nicbVBNS8NAEN34WetXqkcvwSJ4KkkV9Fj04rGC/YA2hM120i7dfLA7UUvMT/HiQRGv/hJv/hu3bQ7a+mDg8d4MM/P8RHCFtv1trKyurW9slrbK2zu7e/tm5aCt4lQyaLFYxLLrUwWCR9BCjgK6iQQa+gI6/vh66nfuQSoeR3c4ScAN6TDiAWcUteSZlX7CvayP8IhZCEjz3DOrds2ewVomTkGqpEDTM7/6g5ilIUTIBFWq59gJuhmVyJmAvNxPFSSUjekQeppGNATlZrPTc+tEKwMriKWuCK2Z+nsio6FSk9DXnSHFkVr0puJ/Xi/F4NLNeJSkCBGbLwpSYWFsTXOwBlwCQzHRhDLJ9a0WG1FJGeq0yjoEZ/HlZdKu15yzWv32vNq4KuIokSNyTE6JQy5Ig9yQJmkRRh7IM3klb8aT8WK8Gx/z1hWjmDkkf2B8/gBBmJSl</latexit>
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- tation behavior (bottom row).
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.
In doing so, the agent is able to 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 (DiAMetR), using learned generative models to imagine new task distributions close to the provided training tasks that can be used to train robust meta-policies.
The contribution of this paper is to propose a framework for making meta-reinforcement learning resilient to a variety of task distribution shifts, and DiAMetR, a practical algorithm instantiating the framework.
DiAMetR trains a population of meta-policies to be robust to different degrees of distribution shifts and then adaptively chooses a meta-policy to deploy based on the inferred test-time distribution shift.
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 memory-based [6, 24], gradient-based [9, 34, 11] and latent-variable based [32, 46, 45] 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 [26, 45, 17] to computer science education [42].
This line of work has been extended to operate in scenarios without requiring any pre-specified task distribution [10, 14] or in offline settings [5, 27, 25] 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 [22, 19] 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 lowest worst-case empirical risk among an “uncertainty-set" of distributions that are boundedly close to the empirical training distribution [36, 20, 2, 13].
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 [36, 20, 2, 13], finding applications in problems like federated learning [13] and image classification [20].
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 [40, 44, 31, 30].
[43] 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]
, T ) = J ( π ∗T ) − E
a ( i ) t ∼ π meta ( ·| h ( i ) t ) , T (cid:34)k k (cid:88) i = T (cid:88) t = r ( i ) t (cid:35) , where h ( i ) t = ( s ( i ) 1: t , r ( i ) 1: t , a ( i ) 1: t − 1 ) ∪ ( s ( j ) 1: T , r ( j ) 1: T , a ( j ) 1: T ) i − 1 j =1 .
Intuitively, the meta-policy has two components: an exploration mechanism that ensures that appropriate reward signal is found for all tasks in the training distribution, and an adaptation 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[9, 32], or implicitly as a black-box RNN [6, 46].
We use the terminology “meta-policies" interchangeably with that of “fast-adaptation" algorithms, since our practical implementation builds on [29] (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 [36] 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 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.
[Table]
[Table]
+JlAZKjQJfdwYUB2rWG4v
+deysM4QQjZdFEvERZG1jgUq8slMBQjTSiTXN9qsQGVlKGOrqhDcGZfnieNStk5KVduTkvVizyOAtknh+SYOOSMVMkVqZE6YeSRPJNX8mY8GS/Gu/ExbV0w8pk98gfG5w+965p/</latexit>
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 [33] or a Wasserstein distance [39]).
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 [31] 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 } M i =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 A), 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 [41]) is
[Figure]
[Figure]
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 [10, 18], 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
⊂ 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 , ϵ 2 )
= (cid:113) ϵ 2 − 1 − 1 ϵ 1 − 1 − 1 , a measure of the mismatch between ϵ 1 and ϵ 2 .
We first compare robust and non-robust solutions by analyzing the bound when ϵ 2 = 0 .
In the regime of β ≪ 1 , excess regret scales as O ( ϵ 1 (cid:113) 1 β ) , meaning that the robust solution is most necessary 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 .
π
meta ,θ , D Replay-Buffer
Meta policy π meta ,θ off
2 [29] θ := θ + α ∇ θ E z ∼ q ϕ
( z ) ( E π ϵ meta ,θ , P (1 k k (cid:88) i =
T (cid:88) t = r ω ( s ( i ) t , a ( i ) t , z ))) 8:
Adversarial task distribution: Update q ϕ using Reinforce [37] ϕ :
= ϕ − α ∇ ϕ
( E z ∼ q ϕ
( z ) [ E π ϵ meta ,θ , P [1 k k (cid:88) i =
T (cid:88) t = r ω ( s ( i ) t , a ( i ) t , z )]] + λD KL ( p train ( z ) ∥
q ϕ ( z )) 9:
Lagrange constraint multiplier: Update λ to enforce D KL ( p train ( z ) ∥
q ϕ ( z )) <
ϵ , λ := λ ≥ 0 λ
+ α ( D KL ( p train ( z ) ∥
q ϕ ( z ))
− ϵ ) 10: for 11: end for
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, and also how the robust meta-RL problem can be solved with stochastic gradient methods.
For simplicity, in the remainder of the paper, we focus on the setting where tasks share transition dynamics, but have different reward functions.
We first introduce the individual components of task parameterization and robust optimization, and describe the overall algorithm in Algorithm 1 and 2.
Parameterizing Task Distributions: Since we assume that variations in tasks correspond to changes in the reward function, the problem of representing a task distribution reduces to learning distributions over reward functions.
We propose to learn a probabilistic model of the task reward functions seen in 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 reward encoder q ψ ( z | h ) that encodes reward samples from an environment history into the latent space, and a decoder r ω ( s, a, z ) mapping a latent vector z to a reward function using a dataset of trajectories collected from the training tasks.
This generative model over reward functions 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]
[Table]
train
[Table]
i
[Table]
[Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] Having learned a latent space, we can parameterize new task distributions q ( T ) as distributions q ϕ ( z ) (the original training distribution corresponds to p train ( z )
=
N (0 , I ) , and measure the divergence between task distributions as well using the KL divergence in this latent space D ( p train ( z ) ∥
q ϕ ( z )) .
Learning Robust Meta-Policies: Given this task parameterization, the next question becomes how 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 [29].
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 [29] as a base learner), and the task proposal distribution with an constrained optimization method (we use dual gradient descent [28]).
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 ( z ) 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 ( z ) satsifies the divergence constraint.
Algorithm
1: Given: p test ( T ) , Π = { π ϵ i meta ,θ } Mi = 2: Initialize TS = Thompson-Sampler() 3: for meta-episode n = 1 , 2 , ... do 4:
Choose meta-policy i =
TS.sample() 5:
Run π ϵ i meta ,θ for meta-episode 6:
TS.update(arm=
i, reward=meta-episode return) 7: end for
Test-time meta-policy selection: Since test-time meta-policy selection can be framed as a multi-armed bandit problem, we use a generic Thompson’s sampling [41] algorithm (see Algorithm 2).
Each metaepisode 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
[Figure]
We aim to comprehensively evaluate DiAMetR and answer the following questions: (1) Do metapolicies learned via DiAMetR allow for quick adaptation under different distribution shifts in 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 train DiAMetR on four continuous control environments: Wheeled navigation [11] (Wheeled driving a differential drive robot), Ant navigation (Ant controlling a four legged robotic quadruped), Fetch reach and Block push [11] (Figures 4a to 4d) (see Appendix H for more details).
Each environment has a train task distribution T i ∼ p
train ( T ) such that each task T i parameterizes a reward function r i ( s, a ) :
= r ( s, a, T i ) .
T i itself remains unobserved, the agent simply has access to reward values and executing actions in the environment.
The meta-policies are evaluated on train task distribution p train ( T ) and on different distributionally shifted test task distribution { p i test ( T ) } Ki =1 .
We use 4 random seeds for all our experiments and include the standard error bars in our plots.
In all of these problems, the distribution of train and test tasks is determined by the distribution of the underlying target locations s t , which determines the reward function (exact distributions in Table 1).
Since these environments have sparse rewards, DiAMetR uses a structured
VAE to model reward distributions (see Appendix C for more details).
6.1 Adaptation to Varying Levels of Distribution Shift
During meta test, given a test task distribution p test ( T ) , DiAMetR uses Thompson sampling to select the appropriate meta-policy π ϵ meta ,θ within N = 250 meta episodes.
π ϵ meta ,θ can then solve any task
T ∼ p test ( T ) within 1 meta episode ( k = 2 environment episode).
Since DiAMetR adaptively chooses a meta-policy during test time, we compare it to RL 2 with test time finetuning.
Figureshows that RL 2 ’s performance more or less remains the same after test time finetuning showing that 10 iteration (with 25 meta-episodes per iteration) isn’t enough for RL 2 to learn an meta-policy for a new task distribution.
For comparison, RL 2 takes 1500 iterations (with 25 meta-episodes per iteration) during training to learn a meta-policy for train task distribution.
To test DiAMetR’s ability to adapt to varying levels of distribution shift, we evaluate it on the above mentioned test task distributions.
We compare DiAMetR with meta RL algorithms such as (off-policy) RL 2 [29], VariBAD [46] and HyperX [47].
Figure 6 shows that DiAMetR outperforms
RL 2 , VariBAD and HyperX on test task distributions.
Furthermore, the performance gap between
DiAMetR and other baselines increase as distribution shift between test task distribution and train task distribution increases.
Naturally, the performance of DiAMetR also deteriorates as the distribution shift is increased, but as shown in Fig 6, 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.
Figureshows that DiAMetR either matches or outperforms RL 2 , VariBAD, and HyperX on the train task distribution.
We refer readers to Appendix D for results on point-navigation environment and
Appendix E for ablation studies and further experimental evaluations.
6.2 Analysis of Tasks Proposed by Latent Conditional Uncertainty Sets
We visualize the imagined test reward distribution for various distribution shifts.
Specifically, we create a heatmap of imagined test reward functions.
Figure 7 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 [Figure]
distribution of rewards has uniformly distributed target locations within the red circle.
As clearly seen in Figure 7, as we increase the distribution shifts, the learned reward distribution model imagines more target locations outside the red circle.
6.3 Analysis of Importance of Multiple Uncertainty Sets
DiAMetR meta-learns a family of adaptation policies, each conditioned on different uncertainty set.
As discussed in section 4, selecting a policy conditioned on a large uncertainty set would lead 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 select an uncertainty set during test time.
To validate this phenomenon empirically, we performed an ablation study in Figure 8.
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 discussed the challenge of distribution shift in meta-reinforcement learning and showed how we can build meta-reinforcement learning algorithms that are robust to varying levels of distribution shift.
We show how we can build distributionally “adaptive" reinforcement learning algorithms that can adapt to varying levels of distribution shift, retaining a tradeoff between fast learning and maintaining asymptotic performance.
We then show we can instantiate this algorithm practically by parameterizing uncertainty sets with a learned generative model.
We empirically showed that this allows for learning meta-learners robust to changes in task distribution.
There are several avenues for future work we are keen on exploring, for instance extending adaptive distributional robustness to more complex meta RL tasks, including those with differing transition dynamics.
Another interesting direction would be to develop a more formal theory providing adaptive robustness guarantees in meta-RL problems under these inherent distribution shifts.
A Test time Meta Policy Selection
As discussed in Section 4, to adapt to test time task distribution shifts, we train a family of metapolicies
Π = { π ϵ i meta } to be robust to varying degrees of distribution shifts.
We then choose the appropriate meta-policy during test-time based on the inferred task distribution shift.
In this section, we frame the test-time selection of meta-policy from the family Π as a stochastic multi-arm bandit problem.
Every iteration involves pulling an arm i which corresponds to executing π ϵ i meta for 1 metaepisode ( k environment episodes) on a task T ∼
p test ( T ) .
Let R i be the expected return for pulling arm
i
[Equation]
Let R ∗ = max i ∈{ 1 ,...,M } R i and
π ϵ meta be the corresponding meta-policy.
The goal of the stochastic bandit problem is to pull arms i 1 , . .
, i N ∈ { 1 , . . .
, M } such that the test-time regret R N is minimized
[Equation]
with constraint that i t can depend only on the information available prior to iteration t .
We choose
Thompson sampling, a zero-regret bandit algorithm, to solve this problem.
In principle, Thompson sampling should learn to choose π ϵ meta after N iterations.
B DiAMetR with Generative Models
To learn a robust meta-policy, we need to define uncertainty sets which are parameterized via q ( T ) and the probability divergence D ( ·||· ) .
In Section 5, we parameterize task (i.e. reward function) distribution as latent space distributions q ϕ ( z ) and measure the divergence between train and test task distribution via KL divergence in latent space D ( p train ( z ) ||
q ϕ ( z )) .
In this section, we show the distributionally robust meta reinforcement learning objective (eq 3) resulting from the abovementioned task distribution parameterization:
[Equation]
This objective function is solved in Algorithm 1 for different values of ϵ .
C Structured VAE for modelling reward distributions
In Section 5, we show how to learn reward function distribution with a variational autoencoder (VAE) [16].
In this section, we leverage the sparsity in reward functions (i.e. 0 / 1 rewards) in the environments used and describe a structured VAE to model r ω ( s, a, z ) with p ( z ) =
N (0 , I ) and
KL-divergence for D ( ·||· ) .
Let h
=
( (cid:80) Tt =1 r t s t ) /
( (cid:80) Tt =1 r t ) be the mean of states achieving a + reward in trajectory h .
The encoder z ∼ q ψ ( z | h ) encodes h into a latent vector z .
The reward model r ω ( s, a, z ) consists of 2 components: (i) latent decoder ˆ h = r hω
( z ) which reconstructs h and (ii) reward predictor r rew ω ( s, ˆ h ) =
exp( −∥ M ⊙ ( s − ˆ h ) ∥ 22 /σ 2 ) which predicts reward for a state given the decoded latent vector.
M is a masking function and σ is a learned parameter.
The training objective becomes
[Equation]
The structure in the VAE helps in extrapolating reward functions when z ∼ q ϕ ( z ) .
This can be further verified by reduction in DiAMetR’s performance on test-task distributions when using vanilla VAE (see Figure 9).
D Experimental Evaluation on Point Robot Navigation
In Section we evaluated DiAMetR on Wheeled-navigation , Ant-navigation , Fetch-reach and Block-push .
We continue the experimental evaluation of DiAMetR in this section and compare it to RL 2 , VariBAD and HyperX on train task distribution and different test task distributions of
Point-navigation environment.
We see that DiAMetR either matches or outperforms the baselines on train task distribution and outperforms the baselines on test task distributions.
Furthermore, adaptively selecting 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.
E Ablation studies
Can improved meta-exploration help a meta-RL algorithm achieve robustness to test-time task distribution shifts?
To test if improved meta-exploration can help meta-RL algorithm achieve robustness to test-time task distribution shifts, we test HyperX [47] on test-task distributions in different environments.
HyperX leverages curiosity-driven exploration to visit novel states for improved meta-exploration during meta-training.
Despite improved meta-exploration, HyperX fails to adapt to test-time task distribution shifts (see Figure 10 for results on Point-navigation and
Figure 6 for results on other environments).
This is because HyperX aims to minimize regret on train-task distribution and doesn’t leverage the visited novel states to learn new behaviors helpful in adapting to test-time task distribution shifts.
Furthermore, we note that the contributions of HyperX is complementary to our contributions as improved meta-exploration would help us better learn robust meta-policies.
Can RL quickly solve test time tasks?
To test if RL can quickly solve test-time tasks, we train
Soft Actor Critic (SAC) [12] on 5 tasks sampled from a particular test task distribution.
To make the comparison fair, we include a baseline that pre-trains SAC on train-task distribution.
Figure 11 shows results on Point-navigation .
We see that both SAC trained from scratch and SAC pre-trained on training task distribution take more than a million timesteps to solve test tasks.
In comparison,
DiAMetR takes 30 k timesteps to select the right meta-policy which then solves new tasks from test distribution in k − 1 environment episodes (i.e. 60 timesteps given k = 2 and horizon H = 60 ).
This shows that meta-RL formulation is required for quick-adaptation to test tasks.
F Proof of Proposition 4.
In this section, we prove the proposition in the main text about the excess regret of an ϵ 2 -robust policy under and ϵ 1 -perturbation (restated below). 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:
E q ( T )
[Regret( π ϵ 2 meta , T ) − Regret( π ϵ 1 meta , T )] ≥[-(cid:18) c ( ϵ 1 , ϵ 2 ) + c ( ϵ 1 , ϵ 2 ) −(cid:19) (cid:112) ϵ 1 (1 − ϵ 1 ) |S 
 | ( |S| − S 0 |-]
[Equation]
The scale of regret depends on c ( ϵ 1 , ϵ 2 )
= (cid:113) ϵ 2 − 1 − 1 ϵ 1 − 1 − 1 , a measure of the mismatch between ϵ 1 and ϵ 2 .
Summary of proof: The proof proceeds in three stages: 1) deriving a form for the optimal meta-policy for a fixed task distribution 1) proving that the optimal ϵ -robust meta-policy takes form:
[Equation]
and finally 3) showing that under the task distribution p ( T g )
= (1 − ϵ 1 ) Uniform ( S 0 ) + ϵ 1 Uniform ( S\S 0 ) , the gap in regret takes the form in the proposition.
Proof.
For convenience, denote S 1 = S\S 0 and Regret( π meta , q ( T )) =
E q [Regret( π, T )] .
Furthermore, since the performance of a meta-policy depends only on its final-timestep visitation distribution (and any such distribution is attainable), we directly refer to π ( g ) as the visited goal distribution of the meta-policy π meta .
Recall that the regret of π meta on task T g is given by 1 π ( g ) .
We begin with the following lemma that demonstrates the optimal policy for a fixed task distribution.
Lemma F.1.
The optimal meta-policy π ∗ q = arg min π Regret( π, q ( T g )) for a given task distribution q ( T g ) is given by
[Equation]
The proof of the lemma is similar to the argument in Gupta et al.
[10], Lee et al.
[Equation]
Letting Z
=
(cid:82) g (cid:112) q
( g ) , we can rewrite the optimization problem as minimizing an f -divergence (with f ( x ) = 1 x )
[Equation]
This is minimized when both are equal, i.e. when π ∗ q ( g ) = √ q/Z , concluding the proof.
Lemma F.2.
The optimal ϵ -robust meta-policy π ∗ ϵ = arg
min R ( π, p train , ϵ ) takes form
[Equation]
Define the distribution q ϵ ( T g ) =
(1 − ϵ ) Uniform ( S 0 )+ ϵ Uniform ( S\S 0 ) , which is an ϵ -perturbation of p train under the TV metric.
We note that there are two main cases: 1) if ϵ = 1 − |S 0 | |S| , then q ϵ is uniform over the entire state space, and otherwise 2) it corresponds to uniformly taking ϵ -mass from
S 0 and uniformly distributing it across S 1 .
Using the lemma, we can derive the optimal policy for q ϵ , which we denote π ϵ :
[Equation]
Writing Z [Equation]
′ , we can write this explicitly as
[Equation]
We now show that there exists no other distribution q ′ ( T ) with TV ( p train , q ′ )
≤ ϵ for which
Regret( π ϵ , q ′ ) ≥ Regret
π ϵ , q ϵ ) .
We break this into the two cases for q ϵ : if q ϵ is uniform over all goals, then π ϵ visits all goals equally often, and so incurs the same regret on every task distribution.
The more interesting case is the second: consider any other task distribution q ′ ( g ) , and let q ′ 0 , q ′ 1 be the probabilities of sampling goals in S 0 and S 1 respectively under q ′ :
q ′
=
E g ∼ q ′
[1( g ∈ S 0 )] and q ′ 1 = 1 − q ′ 0 .
The regret of π ϵ on q ′ is given by
[Equation]
[Equation] By construction of ϵ , we have that (cid:113) |S|−|S [-
-]0 | ϵ ≥
[-|S 
-][Equation]
|S 0 | 1 − ϵ , and so this expression is maximized for the largest value of q ′ 1 .
Under a ϵ -perturbation in the TV metric, the maximal value of q 1 is given by β + ϵ =
ϵ :
[Equation]
This is exactly the regret under our chosen task proposal distribution q ϵ ( T ) (which has q 1 = ϵ ) =
Regret( π ϵ , q ϵ ( T )) .
These two steps can be combined to demonstrate that π ϵ is a solution to the robust objective.
Specifically, we have that
[Equation] [Equation]
so, for any other meta-policy π meta , we have
R
[Equation]
This concludes the proof of the lemma.
Finally, to complete the proof of the original proposition, we write down (and simplify) the gap in regret between π ϵ 1 and π ϵ 2 for the task distribution q ϵ 1 (as described above).
We begin by writing down the regret of π ϵ 1 :
[Equation] [Equation]
Next, we write the regret of π ϵ
Regret( π ϵ 2 , q ϵ 1 ( T )) =
[Equation] ϵ 2 (
[-|S 
 | 1 − ϵ 2 + ϵ(cid:115) |S 1 | ϵ 2 ) (27) = ( (cid:112)-][Equation]
[-|S 
-]ϵ
[Equation]
[Equation]
(1 − ϵ 1 ) (cid:115)
|S
| 1 − ϵ 2 + ϵ(cid:115)
|S 1 | ϵ 2 ) (28) =
(1 − ϵ 1 )
|S
| + ϵ 1 |S 1 | + ϵ(cid:115)
|S 0 ||S 1 | (1 − ϵ 2 ) ϵ
+ (1 − ϵ 1 ) (cid:115)
|S 0 ||S 1 | ( ϵ 2 ) 1 − ϵ(29)
=
(1 − ϵ 1 )
|S
| + ϵ 1 |S 1 | + ϵ(cid:115)
|S 0 ||S 1 | (1 − ϵ 2 ) ϵ
+ (1 − ϵ 1 ) (cid:115)
|S 0 ||S 1 | ( ϵ 2 ) 1 − ϵ(30)
= (1 − ϵ 1 )
|S
| + ϵ 1 |S 1 | + (cid:112)
|S
|S 1 | ϵ 1 (1 − ϵ 1 ) (cid:32)(cid:115) ϵ(1 − ϵ 1 ) (1 − ϵ 2 ) ϵ+
(cid:115) (1 − ϵ 1 )
ϵϵ(1 − ϵ 2 ) (cid:33) (31)
Now writing c ( ϵ 1 , ϵ 2 ) =
(cid:113) ϵ 2 − 1 − 1 ϵ 1 − 1 − 1 = (cid:113) 1 − ϵ 2 ϵϵ 1 1 − ϵ= (1 − ϵ 1 )
|S
| + ϵ 1 |S 1 | + (cid:112)
|S
|S 1 | ϵ 1 (1 − ϵ 1 ) (cid:18)
c
( ϵ 1 , ϵ 2 )
+ c ( ϵ 1 , ϵ 2 ) (cid:19) (32) =
Regret( π ϵ 1 , q ϵ 1 ( T )) + (cid:112) |S
|S 1 | ϵ 1 (1 − ϵ 1 ) (cid:18)
c
( ϵ 1 , ϵ 2 )
+ c ( ϵ 1 , ϵ 2 ) −(cid:19) (33)
This concludes the proof of the proposition.
G Meta-agent Selection and Adaptation during Meta-test
In this section, we show that DiAMetR is able adapt to various test task distributions across different environments by selecting an appropriate meta-agent based on the inferred test-time distribution shift and then quickly adapting the meta-agent to different tasks drawn from the test-distribution.
RL 2 ’s performance remains more or less the same after test-time finetuning showing that 10 iteration (with 25 rollouts per iteration) isn’t enough for RL 2 to learn an adaptive policy for a new task distribution.
[Figure] [Figure]
[Figure] [Figure]
[Figure]
H Environment Description
We describe the environments used in Section 6 and Appendix D of the paper:
• {Point,Wheeled,Ant}-navigation : The rewards for each task correspond to reaching an unobserved target location s t .
The agent (i
[Figure]
e. Wheeled, Ant) must explore the environment to find the unobserved target location (Wheeled driving a differential drive robot, Ant controlling a four legged robotic quadruped).
It receives a reward of 1 once it gets within a small δ distance of the target s t , as in [11].
• Fetch reach : Each task corresponds to moving the gripper to an unobserved target location s t .
The Fetch robot must move its gripper around and explore the environment to find the unobserved target location.
Once the gripper gets within a small δ distance of the target s t , it receives a reward of 1 .
• Block push : Each task corresponds to moving the block to an unobserved target location s t .
The robot arm must move the block around and explore the environment to find the unobserved target location.
Once the block gets within a small δ distance of the target s t , it receives a reward of 1 , as in [11].
Furthermore [Table] Table 2 describes the state space S , action space A , episodic horizon H and frameskip for each environment.
[Table]
I Hyperparameters Used
Table 3 describes the hyperparameters used for the structured VAE for learning reward function distribution
[Table]
We use off-policy RL 2 [29] as our base meta-learning algorithm.
We borrow the implementation from https://github.com/twni2016/pomdp-baselines.
We use the hyperparameters from the config file https://github.com/twni2016/pomdp-baselines/blob/main/configs/meta/ant_dir/rnn.yml but found 1500 num-iters was sufficient for convergence of the meta-RL algorithm.
Furthermore, we use 200 num-updates-per-iter .
Our codebase can be found at https://drive.google.com/drive/folders/1KTjst_n0PlR0O7Ez3-WVj0jbgnl1ELD3?usp=sharing.
We parameterize q ϕ ( z ) as a normal distribution N ( µ, σ ) with ϕ = ( µ, σ ) as parameters.
We use
REINFORCE with trust region constraints (i.e. Proximal Policy Optimization [35]) for optimizing q ϕ ( z ) .
We borrow our PPO implementation from the package https://github.com/ikostrikov/pytorcha2c-ppo-acktr-gail and default hyperparameters from https://github.com/ikostrikov/pytorch-a2c-ppoacktr-gail/blob/master/a2c_ppo_acktr/arguments.py.
Table 4 describes the hyperparameters for PPO that we changed.
[Table]
We use off-policy VariBAD [5] implementation from the package https://github.com/twni2016/pomdpbaselines/tree/main/BOReL with their default hyperparameters.
We use HyperX [47] implementation from the package https://github.com/lmzintgraf/hyperx with their default hyperparameters.
To make the comparisons fair, we ensure that the policy and the Q-function in VariBAD and HyperX have same architecture as that in off-policy RL 2 [29].