Unsupervised Task Clustering for Multi-Task Reinforcement Learning

1 INTRODUCTION Imagine we are given an arbitrary set of tasks . We know that dissimilarities and/or contradicting objectives can exist . However , in most settings we can only guess these relationships and how they might affect joint training . Many recent works rely on such human guesses and ( implicitly or explicitly ) limit the generality of their approaches . This can lead to impressive results , either by explicitly modeling the relationships between tasks as in transfer learning ( Zhu et al. , 2020 ) , or by meta learning implicit relations ( Hospedales et al. , 2020 ) . However , in some cases an incorrect similarity assumption can hurt learning performance ( Lazaric , 2012 ) . Our aim with this paper is to provide an easy , straightforward approach to avoid human assumptions on task similarities . An obvious solution is to train a separate policy for each task . However , this leads to a large amount of experience being required to learn the desired behaviors . Therefore , it is desirable to have a single agent and allow the sharing of knowledge between tasks . This is generally known as multi-task learning , a field which has received a large amount of interest in both the supervised learning and reinforcement learning ( RL ) community ( Zhang & Yang , 2017 ) . If tasks are sufficiently similar , a policy that is trained on one task provides a good starting point for another task , and experience from each task will help training in the other tasks . This is known as positive transfer ( Lazaric , 2012 ) . However , if the tasks are sufficiently dissimilar , negative transfer occurs and reusing a pre-trained policy is disadvantageous . It can even lead to a worse performance than simply starting with a random initialization . Here using experience from the other tasks might slow training or even prevent con- vergence to a good policy . Most previous approaches to multi-task learning do not account for problems caused by negative transfer directly and either accept its occurrence or limit their experiments to sufficiently similar tasks . We present a hybrid approach that is helpful in a setting where the task set contains clusters of related tasks , amongst which transfer is helpful . To illustrate the intuition we provide a conceptualized example in Figure 1 . The figure shows a grid world with 12 tasks that can be naturally clustered in 4 clusters . Note however that our approach goes beyond this conceptual ideal and can be beneficial even if the clustering is not perceivable by humans a-priori . Our approach is inspired by the expectation-maximization framework and uses a set of completely separate policies within our agent . We iteratively evaluate the set of policies on all tasks , assign tasks to policies based on their respective performance and train policies on their assigned tasks . This leads to policies naturally specializing to clusters of related tasks , yielding an interpretable decomposition of the full task set . Moreover , we show that our approach can improve the learning speed and final reward in multi-task RL settings . To summarize our contributions : • We propose a general approach inspired by Expectation-Maximization ( EM ) that can find clusters of related tasks in an unsupervised manner during training . • We provide an evaluation on a diverse set of multi-task RL problems that shows the improved sample complexity and reduction in negative transfer in our approach . • We show the importance of meaningful clustering and the sensitivity to the assumed number of clusters in an ablation study 2 RELATED WORK . Expectation-Maximization ( EM ) has previously been used in RL to directly learn a policy . By reformulating RL as an inference problem with a latent variable , it is possible to use EM to find the maximum likelihood solution , corresponding to the optimal policy . We direct the reader to Deisenroth et al . ( 2013 ) for a survey on the topic . Our approach is different : We use an EM-inspired approach to cluster tasks in a multi-task setting and rely on recent RL algorithms to learn the tasks . In supervised learning , the idea of subdividing tasks into related clusters was proposed by Thrun & O ’ Sullivan ( 1996 ) . They use a distance metric based on generalization accuracy to cluster tasks . Another popular idea related to our approach that emerged from supervised learning is the use of a mixture of experts ( Jacobs et al. , 1991 ) . Here , multiple sub-networks are trained together with an input dependent gating network . Jordan & Jacobs ( 1993 ) also proposed an EM algorithm to learn the mixture of experts . While those approaches have been extended to the control setting ( Jacobs & Jordan , 1990 ; 1993 ; Meila & Jordan , 1995 ; Cacciatore & Nowlan , 1993 ; Tang & Hauser , 2019 ) , they rely on an explicit supervision signal . It is not clear how such an approach would work in an RL setting . A variety of other methods have been proposed in the supervised learning literature , for brevity we direct the reader to the survey by Zhang & Yang ( 2017 ) , which provides a good overview of the topic . Our work differs in that we focus on RL , where no labeled data set exists . In RL , task clustering has in the past received attention in works on transfer learning . Carroll & Seppi ( 2005 ) proposed to cluster tasks based on a distance function . They propose distances based on Q-values , reward functions , optimal policies or transfer performance . They propose to use the clustering to guide transfer . Similarly , Mahmud et al . ( 2013 ) propose a method for clustering Markov Decision Processes ( MDPs ) for source task selection . They design a cost function for their chosen transfer method and derive an algorithm to find a clustering that minimizes this cost function . Our approach differs from both in that we do not assume knowledge of the underlying MDPs and corresponding optimal policies . Furthermore , the general nature of our approach allows it to scale to complex tasks , where comparing properties of the full underlying MDPs is not feasible . An earlier approach by Wilson et al . ( 2007 ) developed a hierarchical Bayesian approach for multi-task RL . Their approach uses a Dirichlet process to cluster the distributions from which they sample full MDPs in the hope that the sampled MDP aligns with the task at hand . They then solve the sampled MDP and use the resulting policy to gather data from the environment and refine the posterior distributions for a next iteration . While their method is therefore limited to simple MDPs , our approach can be combined with function approximation and therefore has the potential to scale to MDPs with large or infinite state spaces which can not be solved in closed form . Lazaric & Ghavamzadeh ( 2010 ) use a hierarchical Bayesian approach to infer the parameters of a linear value function and utilize EM to infer a policy . However , as this approach requires the value function to be a linear function of some state representation , this approach is also difficult to scale to larger problems which we look at . Li et al . ( 2009 ) note that believe states in partially observable MDPs can be grouped according to the decision they require . Their model infers the parameters of the corresponding decision state MDP . Their approach scales quadratically with the number of decision states and at least linearly with the number of collected transitions , making it as well difficult to apply it to complex tasks . More recent related research on multi-task RL can be split into two categories : works that focus on very similar tasks with small differences in dynamics and reward , and works that focus on very dissimilar tasks . In the first setting , approaches have been proposed that condition the policy on task characteristics identified during execution . Lee et al . ( 2020 ) use model-based RL and a learned embedding over the local dynamics as additional input to their model . Yang et al . ( 2020 ) train two policies , one that behaves in a way that allows the easy identification of the environment dynamics and another policy that uses an embedding over the transitions generated by the first as additional input . Zintgraf et al . ( 2020 ) train an embedding over the dynamics that accounts for uncertainty over the current task during execution and condition their policy on it . Our approach is more general than these methods as our assumption on task similarity is weaker . In the second group of papers , the set of tasks is more diverse . Most approaches here are searching for a way to reuse representations from one task in the others . Riemer et al . ( 2018 ) present an approach to learn hierarchical options , and use it to train an agent on 21 Atari tasks . They use the common NatureDQN network ( Mnih et al. , 2015 ) with separate final layers for option selection policies , as well as separate output layers for each task to account for the different action spaces . Eramo et al . ( 2020 ) show how a shared representation can speed up training . They then use a network strucuture with separate heads for each task , but shared hidden layers . Our multi-head baseline is based on these works . Bräm et al . ( 2019 ) propose a method that addresses negative transfer between multiple tasks by learning an attention mechanism over multiple sub-networks , similar to a mixture of expert . However , as all tasks yield experience for one overarching network , their approach still suffers from interference between tasks . We limit this interference by completely separating policies . Wang et al . ( 2020 ) address the problem of openended learning in RL by iteratively generating new environments . Similar to us , they use policy rankings as a measure of difference between tasks . However , they use this ranking as a measure of novelty to find new tasks , addressing a very different problem . Hessel et al . ( 2019 ) present PopArt for multi-task deep RL . They address the issue that different tasks may have significantly different reward scales . Sharma et al . ( 2018 ) look into active learning for multi-task RL on Atari tasks . They show that uniformly sampling new tasks is suboptimal and propose different sampling techniques . Yu et al . ( 2020 ) propose Gradient Surgery , a way of projecting the gradients from different tasks to avoid interference . These last three approaches are orthogonal to our work and can be combined with EM-clustering . We see this as an interesting direction for future work . Quality-Diversity ( QD ) algorithms ( Pugh et al. , 2016 ; Cully & Demiris , 2018 ) in genetic algorithms research aim to find a diverse set of good solutions for a given problem . One proposed benefit of QD is that it can overcome local optima by using the solutions as ” stepping stones ” towards a global optimum . Relatedly in RL , Eysenbach et al . ( 2018 ) and Achiam et al . ( 2018 ) also first identify diverse skills and then use the learned skills to solve a given task . While we do not explicitly encourage diversity in our approach , our approach is related in that our training leads to multiple good performing , distinct policies trained on distinct tasks . This can lead to a policy trained on one task becoming the best on a task that it was not trained on , similar to the ” stepping stones ” in QD . However , in our work this is more a side-effect than the proposed functionality .

Summary: This paper proposes a multi-task RL algorithm that leverages unsupervised task clustering. The authors propose to initialize a number of policies, cluster each task based on its performance on different policies, and train each policy with data coming from tasks within the cluster. The paper shows that such kind of an EM style clustering can lead to better performance than single-task training and be more sample efficient more training each task independently on both tabular settings, continuous control experiments, and Atari.