Reinforcement Learning with Competitive  Ensembles of Information-Constrained Primitives

1 INTRODUCTION . Learning policies that generalize to new environments or tasks is a fundamental challenge in reinforcement learning . While deep reinforcement learning has enabled training powerful policies , which outperform humans on specific , well-defined tasks ( Mnih et al. , 2015 ) , their performance often diminishes when the properties of the environment or the task change to regimes not encountered during training . This is in stark contrast to how humans learn , plan , and act : humans can seamlessly switch between different aspects of a task , transfer knowledge to new tasks from remotely related but essentially distinct prior experience , and combine primitives ( or skills ) used for distinct aspects of different tasks in meaningful ways to solve new problems . A hypothesis hinting at the reasons for this discrepancy is that the world is inherently compositional , such that its features can be described by compositions of small sets of primitive mechanisms ( Parascandolo et al. , 2017 ) . Since humans seem to benefit from learning skills and learning to combine skills , it might be a useful inductive bias for the learning models as well . This is addressed to some extent by the hierarchical reinforcement learning ( HRL ) methods , which focus on learning representations at multiple spatial and temporal scales , thus enabling better exploration strategies and improved generalization performance ( Dayan & Hinton , 1993 ; Sutton et al. , 1999b ; Dietterich , 2000 ; Kulkarni et al. , 2016 ) . However , hierarchical approaches rely on some form of learned high-level controller , which decides when to activate different components in the hierarchy . While low-level sub-policies can specialize to smaller portions of the state space , the top-level controller ( or master policy ) needs to know how to deal with any given state . That is , it should provide optimal behavior for the entire accessible state space . As the master policy is trained 1 Mila , University of Montreal ; 2 Facebook AI Research ; work done while the author was at Mila , University of Montreal ; 3 University of California , Berkeley . anirudhgoyal9119 @ gmail.com on a particular state distribution , learning it in a way that generalizes to new environments effectively becomes the bottleneck for such approaches ( Sasha Vezhnevets et al. , 2017 ; Andreas et al. , 2017 ) . Action selection . Competition mechanism . We argue , and empirically show , that in order to achieve better generalization , the interaction between the low-level primitives and the selection thereof should itself be performed without requiring a single centralized network that understands the entire state space . We , therefore , propose a decentralized approach as an alternative to standard HRL , where we only learn a set of low-level primitives without learning an explicit high-level controller . In particular , we construct a factorized representation of the policy by learning simple primitive policies , which focus on distinct regions of the state space . Rather than being gated by a single meta-policy , the primitives directly compete with one another to determine which one should be active at any given time , based on the degree to which their state encoders “ recognize ” the current state input . While , technically , the competition between primitives implicitly realizes a global selection mechanism , we consider our model decentralized in the sense that individual primitives can function on their own , and can be combined in new ways , without relying on an explicit high-level controller . We frame the problem as one of information transfer between the current state and a dynamically selected primitive policy . Each policy can , by itself , decide to request information about the current state , and the amount of information requested is used to determine which primitive acts in the current state . Since the amount of state information that a single primitive can access is limited , each primitive is encouraged to use its resources wisely . Constraining the amount of accessible information in this way naturally leads to a decentralized competition and decision mechanism where individual primitives specialize in smaller regions of the state space . We formalize this information-driven objective based on the variational information bottleneck . The resulting set of competing primitives achieves both a meaningful factorization of the policy and an effective decision mechanism for which primitives to use . Importantly , not relying on a centralized meta-policy enables the individual primitive mechanisms can be recombined in a plug-and-play fashion , and the primitives can be transferred seamlessly to new environments . Contributions : In summary , the contributions of our work are as follows : ( 1 ) We propose a method for learning and operating a set of functional primitives in a decentralized way , without requiring an explicit high-level meta-controller to select the active primitives ( see Fig . 1 for illustration ) . ( 2 ) We introduce an information-theoretic objective , the effects of which are twofold : a ) it leads to the specialization of individual primitives to distinct regions of the state space , and b ) it enables a competition mechanism , which is used to select active primitives in a decentralized manner . ( 3 ) We demonstrate the superior transfer learning performance of our model , which is due to the flexibility of the proposed framework regarding the dynamic addition , removal , and recombination of primitives . Decentralized primitives can be successfully transferred to larger or previously unseen environments , and outperform models with an explicit meta-controller for primitive selection . 2 PRELIMINARIES . We consider a Markov decision process ( MDP ) defined by the tuple ( S , A , P , r , γ ) , where the state space S and the action space A may be discrete or continuous . The environment emits a bounded reward r : S × A → [ rmin , rmax ] on each transition and γ ∈ [ 0 , 1 ) is the discount factor . π ( .|s ) denotes a policy over the actions given the current state s. R ( π ) = Eπ [ ∑ t γ tr ( st ) ] denotes the expected total return when an agent follows the policy π . The standard objective in reinforcement learning is to maximize the expected total return R ( π ) . We use the concept of the information bottleneck ( Tishby et al. , 2000 ) to learn compressed representations . The information bottleneck objective is formalized as minimizing the mutual information of a bottleneck representation layer with the input while maximizing its mutual information with the corresponding output . This type of input compression has been shown to improve generalization ( Achille & Soatto , 2016 ; Alemi et al. , 2016 ) . 3 INFORMATION-THEORETIC LEARNING OF DISTINCT PRIMITIVES . Our goal is to learn a policy , composed of multiple primitive sub-policies , to maximize the expected reward over T -step interactions for a distribution of tasks . Simple primitives which focus on solving a part of the given task ( and not the complete task ) should generalize more effectively , as they can be applied to similar aspects of different tasks ( subtasks ) even if the overall objective of the tasks are drastically different . Learning primitives in this way can also be viewed as learning a factorized representation of a policy , which is composed of several independent policies . Our proposed approach consists of three mechanisms : 1 ) a mechanism for restricting a particular primitive to a subset of the state space ; 2 ) a competition mechanism between primitives to select the most effective primitive for a given state ; 3 ) a regularization mechanism to improve the generalization performance of the policy as a whole . We consider experiments with both fixed and variable sets of primitives and show that our method allows for primitives to be added or removed during training , or recombined in new ways . Each primitive is represented by a differentiable , parameterized function approximator , such as a neural network . 3.1 PRIMITIVES WITH AN INFORMATION BOTTLENECK . To encourage each primitive to encode information from a particular part of state space , we limit the amount of information each primitive can access from the state . In particular , each primitive has an information bottleneck with respect to the input state , preventing it from using all the information from the state . We define the overall policy as a mixture of primitives , π ( a | s ) = ∑ k ckπ k ( a | s ) , where πk ( a | s ) denotes the kth primitive and ck = δkk′ for k′ ∼ p ( k′ | s ) . We denote the probability of selecting the kth primitive as αk ( s ) : = p ( k | s ) . Rather than learning an explicit model for p ( k | s ) , however , we impose an information-based mechanism for selecting primitives , wherein we limit the amount of information each primitive can contain and select the ones that request the most information about the state . To implement an information bottleneck , we design each of the K primitives to be composed of an encoder penc ( zk | s ) and a decoder pdec ( a | zk ) , together forming the primitive policy , πkθ ( a | s ) = ∫ z penc ( zk | s ) pdec ( a | zk ) dzk . The encoder output zk is meant to represent the information about the current state s that an individual primitive k believes is important to access in order to perform well . The decoder takes this encoded information and produces a distribution over the actions a . Following the variational information bottleneck objective ( Alemi et al. , 2016 ) , we penalize the KL divergence of penc ( zk|s ) and a prior p ( z ) , Lk = DKL ( penc ( zk | s ) ||p ( z ) ) . ( 1 ) In practice , we estimate the marginalization over z using a single sample throughout our experiments . In other words , a primitive pays an “ information cost ” proportional toLk for accessing the information about the current state . In the experiments below , we fix the prior to be a unit Gaussian . In the general case , we can learn the prior as well and include its parameters in θ . The information bottleneck encourages each primitive to limit its knowledge about the current state , but it will not prevent multiple primitives from specializing to similar parts of the state space . To mitigate this redundancy , and to make individual primitives focus on different regions of the state space , we introduce an information-based competition mechanism to encourage diversity among the primitives . 3.2 COMPETING INFORMATION-CONSTRAINED PRIMITIVES . We can use the information measure from equation 1 to define a selection mechanism for the primitives without having to learn a centralized meta-policy . The intuition is that the information content of an individual primitive encodes its effectiveness in a given state s such that the primitive with the highest value Lk should be activated in that particular state . In particular , we set αk = Z−1 exp ( βLk ) to obtain a distribution over k as a function of the information content , activating the primitives with the highest information content . Here , Z =∑ k exp ( βLk ) is a normalization constant . This mechanism enables competition between primitives , leading them to focus on parts of the state space that they “ understand ” well and letting others act in other parts . Trading reward and information . To perform proper credit assignment , the environment reward is distributed to primitives according to their participation in the global decision , i.e . the reward rk given to the kth primitive is weighted by its selection coefficient , such that rk = αkr , with r = ∑ k rk . Hence , a primitive can potentially get a higher reward when deciding to act , but it also pays a higher price for accessing more information about the current state . The information bottleneck and the competition mechanism , when combined with the overall reward maximization objective , will lead to specialization of individual primitives to distinct regions in the state space . That is , each primitive should specialize in a part of the state space that it can reliably associate rewards with . Since the entire ensemble still needs to understand all of the state space for the given task , different primitives need to encode and focus on different parts of the state space .

Summary: The paper draws upon the idea of information bottleneck to do task decomposition so as to learn policy primitives similar to hierarchical reinforcement learning that combine together in a competitive manner to specialize in different parts of the task's domain. These policy primitives don't need a higher-level meta-policy to stitch them together. Instead the decision is made in a decentralized manner balancing the cost of information acquisition with maximizing rewards.