Taming the Herd: Multi-Modal Meta-Learning with a Population of Agents
Abstract
Meta-learning is a paradigm whereby an agent is trained with the speciﬁc goal of fast adaptation.
With promising recent results, it is an effective way to use prior knowledge for future learning.
However, most of the prominent successes focus entirely on adapting to similar tasks, from a unimodal distribution.
This drastically differs from the real world, where problems may be more diverse.
In this paper we address this issue, and provide a simple approach to solve it, which we call Taming the Herd .
The Herd refers to a population of agents, each specializing on a subset (or mode) of the task distribution.
At test time, we automatically allocate the appropriate member of the Herd and thus perform comparably with an oracle trained solely on those tasks.
We apply our approach to both MAML and PEARL, and demonstrate its efﬁcacy on a simple yet challenging multi-modal task distribution.
1. Introduction
In recent times there has been great excitement in the ﬁeld of meta-learning, where agents are trained with the goal of adapting to new tasks.
This approach resonates with many since it offers the promise of agents which can learn something more fundamental than just an individual task.
Our focus is on meta- reinforcement -learning, where agents can quickly learn new RL tasks
Model-agnostic meta-learning (MAML, (Finn et al., 2017)) is a popular approach, which seeks to learn a set of parameters which can successfully adapt to a new task in just a few steps.
Its simplicity and effectiveness make it a canonical algorithm in the space, inspiring a series of follow-up work.
One crucial limitation of MAML is its requirement to train on unimodal task distributions.
This means that in reality, it can only adapt to a few very similar tasks.
Thus far, existing approaches to improve MAML have focused on a single dimension, improving efﬁciency (Rakelly et al., 2019;
Behl et al., 2019).
This means that despite vast improvements over the past few years, none of the latest algorithms perform well when presented with multimodal task distributions.
In this paper we take the ﬁrst step towards solving this problem.
We use task embeddings to learn clusters of similar tasks, in an unsupervised manner and train a population (or Herd) of meta-learners, each specialized to one of the learned task clusters.
At test time we use a probe policy to assign the new task to the best member of the learned cluster population, which can then rapidly adapt in the remaining test steps.
We call our approach: Taming the Herd and depict it in ﬁgure 1.
Our contributions are twofold: ﬁrst we introduce a mechanism to meaningfully measure similarity between tasks for the purpose of meta-learning.
To do this, we introduce a variety of task embeddings.
Second, we introduce a way to leverage the task geometry induced by these task embeddings to improve existing meta-learning algorithms such as
MAML (Finn et al., 2017) and PEARL (Rakelly et al., 2019).
We show that by incorporating these insights, both MAML and PEARL can deal with multi-modal task distributions, where they would otherwise fail.
The paper is organized as follows.
We begin with related work, before introducing RL and the meta-learning problem in Section 3.
Next we introduce our algorithm, before demonstrating the efﬁcacy of our approach on a challenging multi-modal task.
2. Related Work
Meta-learning as a concept has been around for decades (Schmidhuber, 1987; Naik and Mammone, 1992).
However, with the introduction of model-agnostic meta-learning (MAML (Finn et al., 2017)), it has experienced a surge of interest.
MAML seeks to learn a set of weights which are amenable to fast adaptation, by backpropagating through the update step.
This has led to a series of follow up works (Antoniou et al., 2019; Behl et al., 2019; Rajeswaran et al., 2019), all of which assume access to a distribution of related tasks.
We will introduce MAML more formally in Sec: 3.
More recently, a series of works have considering alternative ways to explicitly decouple the task inference module, from the learned shared representation that allows the agent to generalize across tasks.
One existing approach is to explicitly allocate a subset, or a low dimensional representation of the parameters to be updated during the adaptation phase such as (Zintgraf et al., 2019; Rusu et al., 2019).
A competing approach, and one we explore further in this work is to learn a task embedding aimed at inferring the nature of the task at hand and produce a policy conditioned on this information (Rakelly et al., 2019; Zintgraf et al., 2020).
Despite these advances, these approaches fail to deal with multimodal task distributions.
We show that with our approach,
PEARL can perform well in this setting.
Our work is connected to Task2Vec
(Achille et al., 2019) in that we also embed tasks.
While their work selects the best classiﬁer on a given task we do not stop at this point but use the identiﬁed closest ensemble member for meta-testing.
In addition, another recent paper discussed the reasons for
MAMLs effectiveness, suggesting this is partially due to effective feature learning (Raghu et al., 2020).
The performance gap for a single MAML suggests that in fact these features may not generalize across multiple modes of a task distribution.
The problem of multi-modal learning is not unique to meta-learning.
Existing solutions to similar problems arising in model based RL (see (Kaushik et al., 2020)) or natural language processing (Vu et al., 2020) bear resemblance to our work.
3. Background 3.1. Reinforcement Learning
A Markov Decision Process (MDP, (Bellman, 1957)) is a tuple ( S , A , P , R) .
Here S and A stand for the sets of states and actions respectively, such that for s t , s t +1 ∈ S and a t ∈ A : P( s t +1 | s t , a t ) is the probability that the system/agent transitions from s t to s t +1 given action a t and
R( a t , s t , s t +1 ) is a reward obtained by an agent transitioning from s t to s t +1 via a t .
A policy π θ : S → A is a (possibly randomized) mapping (parameterized by θ ∈ R d , e.g. weight of the neural network) from S to A .
Policy learning is the task of optimizing parameters θ of π θ such that an agent applying it in the environment given by a ﬁxed MDP maximizes total (expected/discounted) reward over given horizon H .
In this paper we consider MDPs with ﬁnite horizons.
We denote the trajectory of a policy as τ , deﬁned as the sequence of states and actions taken during an episode, τ = { s 0 , a 0 , r 0 , . . .
s H , a H , r H } .
3.2. Meta-Learning
When moving to meta-RL, we ﬁrst deﬁne a task distribution p ( T ) .
Here each task is an MDP as described.
We denote T i as a task sampled from p ( T ) , such that
T i = { p ( s 0 ) , P i ( s t +1 | s t , a t ) , R i ( a t , s t , s t +1 ) } .
As such, the loss function for a given task is generally of the following form:
[Equation]
Meta-learning considers the problem of fast adaptation .
In
MAML this means that we no longer consider a loss function with respect to the parameters θ , but
those θ (cid:48) i , where:
[Equation]
for some learning rate α .
Thus, the meta-RL loss function is as follows:
[Equation]
) in other words, we consider the loss on each task after taking one gradient step on the task.
In practice, we also consider increasing beyond a single gradient step.
Note the crucial component here, that the tasks T i are sampled iid from p ( T ) .
This means that all existing approaches can only work when trained and tested on similar tasks, reducing the ability to learn multiple different tasks, as present in the real world.
4. Taming the Herd
We now present our method: Taming the Herd.
We begin by training a single meta-learning policy across all tasks.
This policy is referred to as the probe policy, π θ probe .
We then use this policy to gain experience from each task, used to form a task embedding, which we denote φ ( T i ) .
Once we have the embeddings (we use φ ( T ) to refer to the embeddings for all tasks), we cluster the tasks into K clusters
C K , which we refer to as Cluster .
We typically used hierarchical clustering for our experiments, but other methods could be considered.
Once we have the K clusters, we train a population of K meta-learning policies { θ k } k ∈ K on each cluster of tasks in C k .
We refer to this procedure as
MetaLearn , and this could be any meta-learning algorithm, but in practice we apply our approach to MAML (Finn et al., 2017) and PEARL (Rakelly et al., 2019).
The full meta-train procedure is shown in Algorithm 1.
Initialize
[Equation]
At test time, we ﬁrst use the probe policy to gather sufﬁcient experience (typically one episode).
We then use this data to match the task with the corresponding cluster.
At this point, we can use the remaining meta-test gradient steps to update the corresponding member of the population on the task.
The full meta-test procedure is shown in Algorithm 2.
i ) embedding i )) where c i is the centroid of the i th cluster.
We now use the corresponding policy θ i
Meta-Test: MetaTest( θ i , T test )
When applied to MAML we refer to our method as a Herd of
MAMLs (HoM), while for PEARL, we introduce a String of PEARLs (SoP).
Next we go into more detail on each component of our method.
4.1. Probe policies
A crucial part of the algorithm is the choice of the probe policy π θ probe .
We wish to obtain task embeddings by rolling out π θ probe , such that modes of tasks are clustered together.
That means a) very similar tasks should be clustered together and b) if there is a task B such that success on A is only possible if jointly trained with B, φ A and φ B are also close together.
It is important that the probe policy can quickly learn behaviors which can then be used to distinguish between tasks.
We focus on a simple multi-task probe policy principle any multi-task method, and use the baseline single version of the MAML or PEARL algorithm, but other multitask methods could be considered (Teh et al., 2017).
We empirically found randomly initialized policy-networks to be insufﬁcient.
The design of speciﬁc acquisition functions for task inference is left to exciting future work.
4.2. Task Embeddings
Depending on the choice of π θ probe there are a multitude of options to ﬁnd an embedding φ of task T i .
We propose several here, with discovery of others left to future work.
For Taming the Herd, we embed each transition individually and average the transition embeddings along a trajectory to get a trajectory embedding.
We propose the following approaches to embed { s t , a t , r t , s t +1 } :
The Fisher information matrix (FIM, deﬁned in the Appendix Sec.
A) can be seen as the curvature of the parameter-space.
Thus we consider a reward weighted FIM embedding
[Equation]
[Equation]
, where ˜ a t ∼ π θ ( ·| s t ) , as we want the Fisher and not the empirical Fisher.
a reward weighted state-next-state embedding as product of reward and the outer product of concatenated state - next state tuples:
[Equation]
When using a PEARL as probe policy a particular natural choice of task-embedding is the inferred latent context of the exploration trajectories of
task
T i 4.
A reward weighted gradient embedding is:
[Equation]
We found empirically that normalizing the rewards across the observed trajectories of each task works best, thus we use r t = r t / max t r t .
We focus on the reward weighted FIM embedding for our experiments, and plot obtained task embeddings in the appendix.
4.3. Number of clusters
One of the key choices of our algorithm is the number of clusters K we assign the tasks to.
Given N tasks any number in the range [1 , N ] is a candidate.
Choosing K = recovers the basic case to train one meta-learner on all test tasks, allowing for positive transfer between all tasks but at the expensive of possibly one task hindering the others.
Choosing K
= N is the other extreme, here one metalearner is trained for each task, allowing for no knowledge transfer between the tasks.
While this might be optimal for very different tasks, most real world tasks share some common struture where the seperate training prevents the exchange of information.
The goal of the task embeddings presented in 4.2 is to embed tasks in a way such that similar tasks are close to each other.
5. Experiments
The primary goal of our experiments it to evaluate whether our approach to meta-learning can improve performance on multi-modal task distributions.
This differs from existing works, which only consider variations of the same task.
We create a multi-modal distribution by taking the union of two common meta-RL problems: { Cheetah : Forward , Cheetah : Backward } and { Walker : Forward , Walker : Backward } .
These two sets of tasks correspond to different agents (a Cheetah and a
Walker), which get a positive reward walking in different directions (forward or backward).
We call this set of tasks
ML2.
We focus our investigation on two questions:
Is it beneﬁcial to train a population of meta-learners instead of a
single one? 2.
Are the task embeddings from Section 4.2 suitable proxies to the task-assignment-oracle?
Our implementation is based on (garage contributors, 2019).
We use the default the default hyperparamaters from garage for MAML, and the default for PEARL from the original paper (Rakelly et al., 2019).
For completeness, we include these in the Appendix.
[Figure]
To ensure a fair comparison of the single meta-learner and the herd we make sure that the single meta-learner has at least as many environment steps at test and train time as all the ensemble members together.
We now describe the training protocol in detail for MAML on ML2.
At train time the probe-policy is trained for x episodes for a total of y steps per task.
Subsequently each member of the herd is trained on z samples form task i, meaning the single meta-learner is trained on i × N total environment interactions.
This means the evaluation protocol differs.
If fewer environment steps than the total number of probe-policy training steps have been taken, we are still in the single task case.
At test time the single meta-learner uses l exploration trajectories to do one-shot learning on top of it.
The herd approach can use only l − 2 adaptation trajectories as we dedicate a samplebudget of two trajectories to embed the test-task with the probe-policy.
For PEARL we embed the task at test-time with only one trajectory and use 1 exploration trajectory.
[Table]
[Table]
We show the training curved in Fig.
2, where we see that in both cases, our approach achieves dramatically better performance than the baseline of a single agent.
In Table 1 & 2 we show the median best performance achieved for the MAML and PEARL experiments respectively.
For MAML, we see that a single MAML outperforms on the HalfCheetah task, but it is signiﬁcantly worse on Walker .
Interestingly, we see larger gains for PEARL than MAML.
6. Discussion
We presented Taming the Herd , a new approach for multimodal meta-learning using a population of agents.
Our results show signiﬁcant promise, dramatically outperforming the baseline of a single agent, for two different popular metalearning algorithms, MAML and PEARL.
We believe our approach achieves positive forward and backward transfer.
For forward transfer, the results in this work demonstrate the performance in few shot learning.
However, potentially more interesting is the concept of backward transfer.
We believe adding more tasks may improve the meta learned representations for the corresponding clusters.
This approach leads to a variety of exciting future directions.
On the individual work includes off-policy training at meta test-time to use the trajectories obtained by the probe policy not only for task-inference but also for meta training.
Advances in this could be applied also to PEARL which does also only task-inference at test-time without adapting the learners parameters.
We can also consider measuring distances between probabilistic embeddings, for example using Wasserstein Distances (Pacchiano et al., 2020).
We could consider other objectives for the probe policy.
Currently, we solely optimize for reward.
It would be interesting to consider training the probe policy such that it explicitly seeks to discover its’ mode.
Measures employing active learning have recently achieved success in RL (Ball et al., 2020; Pathak et al., 2019), and may be useful here.
7. Acknowledgement
RM was supported by appliedAI.
A. Additional Deﬁnitions A.1. Fisher Information Matrix
The Fisher information matrix (FIM) is deﬁned as:
[Equation]
which is the expected covariance of the gradients of the log-likelihoods w.r.t the model parameters.
The information content of each weight of a neural network can be computed by the divergence of the original and a perturbed output distribution for w (cid:48)
=
w + δw in second order approximation as: E x ∼ p KL ( p w (cid:48)
( y | x ); p w ( y | x )) =
δw ˙
Fδw + o ( δw 2 ) .
The FIM can be seen as the curvature of the model in the parameter-space.
B. Implementation Details
We used hyperparameters from open source repos.
For MAML we used garage (garage contributors, 2019) and for PEARL we used the author’s implementation (Rakelly et al., 2019).
For completeness we list these below.
B.1. Hyperparameters MAML + TRPO
[Table]
[Figure]
B.2. Hyperparameters PEARL C. Additional Experimental Results C.1. Plot of all embeddings
[Figure]
[Figure]