Collaborative Inter-agent Knowledge Distillation for Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) [ 47 ] has demonstrated impressive performance on solving sequential decision-making tasks in interactive environments ( e.g. , video game-playing [ 32 ] or robotic control [ 18 ] ) . In these tasks , the RL agent ’ s goal is to find the optimal policy , which maximizes the expected return in the task . The agent learns this optimal policy through many trial-and-error interactions with the environment . Through these trial-and-error interactions , the agent explores the consequences of different decisions , in the form of a reward or punishment . Unfortunately , millions of interactions are required even to solve simple tasks . Such amounts of interaction are infeasible to acquire in real , physical tasks , as is the case for robots . In order to learn the optimal policy , the agent performs policy improvement , which refers to the incremental process of refining the agent ’ s policy towards performing high-rewarding actions . However , if the agent has a poor policy , then this incremental updating can be a slow process , requiring a large number of interactions before the agent acquires the experiences necessary for learning desirable behaviors . This is because a typical RL agent gathers experiences using a policy close to its greedy policy , and if this greedy policy itself is poor , acquiring high-quality experiences can require a large number of exploratory actions . However , through collaboration amongst multiple agents with different behaviors , perhaps we can reduce the amount of experiences needed for an agent to acquire a good policy . For inspiration , we turn to the study of collective animal behavior [ 46 ] . Rather than exclusively learn from trial-and-error alone , the optimal behaviors for several tasks ( e.g. , flocking and foraging ) can emerge through collaboration amongst animals . For example , consider ants foraging in a colony [ 8 ] . The ants search divergent paths , resulting in extensive exploration as a group , across a wide range of food-gathering policies . Furthermore , the ants that find food share their path with their companions via pheromones . This information-sharing amongst colony members offers each ant beneficial guidance toward food , so that individual ants need not exhaustively search several paths in order to find food . To make an analogy to the RL setting , the ants ’ foraging process can be viewed as an effective way of policy improvement . Motivated by these insights from collective animal behavior , we incorporate the notion of collective knowledge sharing into the RL setting in an attempt to accelerate and guide the search for the optimal policy . Our method emulates the collaborative behaviors of ants via an ensemble of RL agents : a group of agents collectively search for the optimal policy in the same task , while periodically sharing knowledge . Each RL agent resembles an individual in the ant colony . To elicit diverse experiences amongst members of the ensemble , each RL is randomly initialized with different neural network parameters . Random initialization results in adequate behavioral diversity [ 22 , 34 ] needed for collective exploration . As these agents are diverse in nature , at any given time during the course of training , one agent naturally has a policy superior to its peers . This agent can guide its peers towards higher-performing policies , just as ants share superior paths through pheromones . In order to guide the other agents , our method employs a knowledge distillation framework [ 24 ] which is effective at transferring knowledge between neural networks , without assuming identical model architectures . Knowledge distillation encourages other agents to act in a manner similar to the leading agent , allowing agents to escape underperforming policies . Moreover , knowledge distillation empowers all agents to continue searching for the optimal policy from better starting points , accelerating knowledge accumulation in the entire ensemble . However , despite the distillation , the agents ’ knowledge is still preserved , retaining diversity amongst the ensemble . This paper ’ s primary contribution is collaborative inter-agent knowledge distillation ( CIKD ) , a simple yet effective framework for RL that jointly trains an ensemble of RL agents while periodically performing knowledge sharing . We demonstrate empirically that our method can improve the state-ofthe-art soft-actor critic ( SAC ) [ 20 ] on a suite of challenging MuJoCo tasks , exhibiting superior sample efficiency and performance . We further validate the effectiveness of distillation for knowledge sharing by comparing against other methods of sharing knowledge . In addition , we present an in-depth investigation to explain the underlying causes of CIKD ’ s performance improvement . Finally , our ablation study shows that a small ensemble is sufficient for improving performance . The remainder of this paper is organized as follows . Section 2 discusses the related work . Section 3 introduces the reinforcement learning formulation . Section 4 describes our method . Section 5 presents our experimental findings . Section 6 summarizes our contributions and outlines potential avenues for future work . 2 RELATED WORK . The idea of jointly training multiple policies through RL has emerged in prior works . The most relevant works [ 37 , 38 ] train multiple policies for the same task through RL , as our method does . Osband et al . [ 37 , 38 ] train several agents in an ensemble while storing these agents ’ experiences in a shared buffer . Thus , agents share knowledge by sharing experiences amongst members of the ensemble , which are then used for RL updates . Our method is complementary to Osband et al . [ 37 , 38 ] ’ s , in that we can also use a shared buffer of experiences , but we additionally periodically performing knowledge distillation between members of the ensemble , in particular from the best agent to other agents . In our experiments we present supporting evidence to justify the importance of knowledge sharing . Other related methods aggregate multiple policies to select actions [ 2 , 6 , 10 , 17 , 31 , 40 , 41 , 45 , 50 , 54 ] . Abel et al . [ 1 ] , Tosatto et al . [ 52 ] , Wang & Jin [ 53 ] sequentially train a series of policies , boosting the learning performance by using the errors of a prior policy . However , rather than perform decision aggregation or sequentially-boosted training , we focus on improving the performance of each individual agent via knowledge sharing amongst jointly trained agents . Again , our method can be considered to be a complementary approach to decision aggregation and sequentially boosted training methods . We can view the sharing of knowledge from the best policy as exploiting successful behavior patterns . This general notion has been explored in several areas of RL . Rusu et al . [ 42 ] and Parisotto et al . [ 39 ] can train a single network to perform multiple tasks by transferring multiple pre-trained RL agents ’ policies to a single network . Hester et al . [ 23 ] , Nair et al . [ 35 ] accelerate the RL agents ’ training progress via human experts ’ guidance . Instead of using experts ’ policies , Levine & Koltun [ 30 ] , Nagabandi et al . [ 33 ] and Zhang et al . [ 55 ] leverage model-based controllers ’ behaviors ( e.g . model predictive controllers [ 15 ] or linear-quadratic regulators Dorato et al . [ 7 ] ) , facilitating training for RL agents . Additionally , Hong et al . [ 25 ] and Oh et al . [ 36 ] train agents to imitate the past successful self-experiences or policies . Orthogonal to the aforementioned works , our method periodically exploits the current best policy amongst the ensemble , and shares it with the ensemble , enabling a collaborative search for the optimal policy . Collaborative learning approaches have emerged in other areas of machine learning research . In computer vision , Zhang et al . [ 56 ] present deep mutual learning , which trains multiple models that mutually imitate each other ’ s outputs on classification tasks . Our distillation is not mutual , and flows in a single direction , from a superior teacher agent to other student agents in the ensemble . In subsequent work by Lan et al . [ 29 ] , they train an ensemble of models to imitate a stronger teacher model . Simultaneously , this teacher model learns to aggregate all of the ensemble models ’ predictions . They demonstrate that imitating a superior teacher leads to better performance than deep mutual learning , which performs mutual imitation amongst members . Our method contrasts from the above methods by periodically electing the teacher for distillation . Since our distillation occurs periodically , we can collect performance statistics for each agent between periods , which can be used to elect the teacher in the distillation phase . In contrast to the distillation of an aggregate policy , as Lan et al . [ 29 ] do , our method ensures that the teacher indeed possesses a superior policy . Collaborative learning has also been explored in RL as well . Teh et al . [ 49 ] and Ghosh et al . [ 16 ] jointly learn independent policies for multiple tasks or contexts and then distill these policies to a central multi-task policy . Galashov et al . [ 12 ] learn a task-specific policy while bounding the divergence between this task-specific policy and some generic policy that can perform basic taskagnostic behaviors . Czarnecki et al . [ 5 ] gradually transfer the knowledge of a simple policy to a complex policy during the course of joint training . While our method also collaboratively trains policies , we differ from the aforementioned works in several aspects . First , our method periodically elects a leading agent for sharing knowledge rather than either constraining the mutual policy divergence [ 14 , 16 , 49 , 56 ] or imitating aggregated models [ 29 ] . The second difference is that our method does not rely on training heterogeneous policies ( e.g . a simple policy and a complex policy [ 5 ] ) , which makes our method more generally applicable . Finally , as opposed to Teh et al . [ 49 ] and Ghosh et al . [ 16 ] , we consider single-task setting rather than multi-tasking . Evolutionary algorithms ( EA ) [ 13 , 21 , 26 , 43 ] similarly employ multiple policies to find the optimal policy . EA repeatedly performs mutation , selection , and reproduction on a maintained population . Mutation randomly perturbs the parameters of policies in the population ; selection eliminates the underperforming policies by testing the policies ’ performance in the environment ; reproduction produces the next generation of policies from the remaining policies . Unlike EA , our method does not rely on mutation and reproduction on a population of policies . While EA and our method are similar in an abstract sense , in that they both share knowledge amongst agents , they are quite different in practice . EA often eliminates members from its population and performs destructive changes to members of the population . Our method focuses on continuously improving the existing agents , and does not perform destructive changes to its population . In fact , our method can be incorporated into EA , serving as a more effective way to optimize each individual policy within the same generation . 3 BACKGROUND . In this section we describe the general framework of RL . RL formalizes a sequential decision-making task as a Markov decision process ( MDP ) [ 47 ] . An MDP consists of a state space S , a set of actions A , a ( potentially stochastic ) transition function T : S ×A → S , a reward functionR : S ×A → R , and a discount factor γ ∈ [ 0 , 1 ] . An RL agent performs episodes of a task where an agent starts in a random initial state s0 , sampled from the initial state distribution ρs0 , and performs actions , experiencing new states and rewards . More generally , at timestep t , an agent performs an action at in state st , receives a reward rt+1 , and transitions to a new state st+1 , according to the transition function T . The discount factor γ is used to indicate the agent ’ s preference for short-term rewards over long-term rewards . An RL agent performs actions according to its policy , a conditional probability distribution πφ : S ×A 7→ [ 0 , 1 ] , where φ denotes the parameters of the policy , which may be the parameters of a neural network . RL methods iteratively update φ via rollouts of experience τ = { ( st , at , rt , st+1 ) } Tt=0 , seeking within the parameter space Φ for the optimal φ∗ that maximizes the expected return Es∼ρs0 [ ∑T t=0 γ trt|s0 = s ] for each t within an episode .

Summary: This paper proposes an RL training procedure that maintains an ensemble of k policies and periodically pushes all the policies to be closer to the best performing one. The formulation, experiments and analysis are very clear and show a mild improvement over using the same underlying RL algorithm without the imitation part. The idea is close to many other proposed in the literature, but to my knowledge it is the first time this exact procedure is studied in detail.