Modular Lifelong Reinforcement Learning via Neural Composition

1 INTRODUCTION . Reinforcement learning ( RL ) has achieved impressive success at complex tasks , from mastering the game of Go ( Silver et al. , 2016 ; 2017 ) to robotics ( Gu et al. , 2017 ; Andrychowicz et al. , 2020 ) . However , this success has been mostly limited to solving a single problem given enormous amounts of experience . In contrast , humans learn to solve a myriad of tasks over their lifetimes , becoming better at solving them and faster at learning them over time . This ability to handle diverse problems stems from our capacity to accumulate , reuse , and recombine perceptual and motor abilities in different manners to handle novel circumstances . In this work , we seek to endow artificial agents with a similar capability to solve RL tasks using functional compositionality of their knowledge . In lifelong RL , the agent faces a sequence of tasks and must strive to transfer knowledge to future tasks and avoid forgetting how to solve earlier tasks . We formulate the novel problem of lifelong RL of functionally compositional tasks , where tasks can be solved by recombining modules of knowledge in different manners . While temporal compositionality has long been studied in RL , such as in the options framework , the type of functional compositionality we study here has not been explored in depth , especially not in the more realistic lifelong learning setting . Functional compositionality involves a decomposition into subproblems , where the outputs of one subproblem become inputs to others . This moves beyond standard temporal composition to functional compositions of layered perceptual and action modules , akin to programming where functions are used in combination to solve different problems . For example , a typical robotic manipulation solution interprets perceptual inputs via a sensor module , devises a path for the robot using a high-level planner , and translates this path into motor controls with a robot driver . Each of these modules can be used in other combinations to handle a variety of problems . We present a new method for continually training deep modular architectures , which enables efficient learning of the underlying compositional structures . The modular lifelong RL agent will encounter a sequence of compositional tasks , and must strive to solve them as quickly as possible . Our proposed solution separates the learning process into three stages . First , the learner discovers how to combine its existing modules to solve the current task to the best of its abilities by interacting with the environment with various module combinations . Next , the agent accumulates additional information about the current task via standard RL training with the optimal module combination . Finally , the learner incorporates any newly discovered knowledge from the current task into existing modules , making them more suitable for future learning . We demonstrate that this separation enables faster training and avoids forgetting , even though the agent is not allowed to revisit earlier tasks for further experience . Our main contributions include : 1 . We formally define the lifelong compositional RL problem as a compositional problem graph , encompassing both zero-shot generalization and fast adaptation to new compositional tasks . 2 . We propose two compositional evaluation domains : a discrete 2-D domain and a realistic robotic manipulation suite , both of which exhibit compositionality at multiple hierarchical levels . 3 . We create the first lifelong RL algorithm for functionally compositional structures and show empirically that it learns meaningful compositional structures in our evaluation domains . While our evaluations focus on explicitly compositional RL tasks , the concept of functional composition is broadly applicable and could be used in future solutions to general lifelong RL . 4 . We propose to use batch RL techniques for avoiding catastrophic forgetting in a lifelong setting and show that this approach is superior to existing lifelong RL methods . 2 RELATED WORK . Lifelong or continual learning Most work in lifelong learning has focused on the supervised setting , and in particular , on avoiding catastrophic forgetting . This has typically been accomplished by imposing data-driven regularization schemes that discourage parameters from deviating far from earlier tasks ’ solutions ( Zenke et al. , 2017 ; Li & Hoiem , 2017 ; Ritter et al. , 2018 ) , or by replaying real ( Lopez-Paz & Ranzato , 2017 ; Nguyen et al. , 2018 ; Chaudhry et al. , 2019 ; Aljundi et al. , 2019 ) or hallucinated ( Achille et al. , 2018 ; Rao et al. , 2019 ; van de Ven et al. , 2020 ) data from earlier tasks during the training of future tasks . Other methods have instead aimed at solving the problem of model saturation by increasing the model capacity ( Yoon et al. , 2018 ; Li et al. , 2019 ; Rajasegaran et al. , 2019 ) . A few works have addressed lifelong RL by following the regularization ( Kirkpatrick et al. , 2017 ) or replay ( Isele & Cosgun , 2018 ; Rolnick et al. , 2019 ) paradigms from the supervised setting , exacerbating the stability-plasticity tension . Others have instead proposed multi-stage processes whereby the agent first transfers existing knowledge to the current task and later incorporates newly obtained knowledge into a shared repository ( Schwarz et al. , 2018 ; Mendez et al. , 2020 ) . We follow the latter strategy for exploiting and subsequently improving accumulated modules over time . Compositional supervised learning While deep nets in principle enable learning arbitrarily complex task relations , monolithic agents struggle to find such complex relations given limited data . Compositional multi-task learning ( MTL ) methods use explicitly modular deep architectures to capture compositional structures that arise in real problems . Such approaches either require the compositional structure to be provided ( Andreas et al. , 2016 ; Hudson & Manning , 2018 ) or automatically discover the structure in a hard ( Rosenbaum et al. , 2018 ; Alet et al. , 2018 ; Chang et al. , 2019 ) or soft ( Kirsch et al. , 2018 ; Meyerson & Miikkulainen , 2018 ) manner . A handful of approaches have been proposed that operate in the lifelong setting , under the assumption that each component can be fully learned by training on a single task and then reused for other tasks ( Reed & de Freitas , 2016 ; Fernando et al. , 2017 ; Valkov et al. , 2018 ) . Unfortunately , this is infeasible if the agent has access to little data per task . Recent work proposed a framework for lifelong supervised compositional learning ( Mendez & Eaton , 2021 ) , similar in high-level structure to our proposed method for RL . Typical evaluations of compositional learning use standard benchmarks such as ImageNet or CIFAR100 . While this enables fair performance comparisons , it fails to give insight about the ability to find meaningful compositional structures . Some notable exceptions exist for evaluating compositional generalization in supervised learning ( Bahdanau et al. , 2018 ; Lake & Baroni , 2018 ; Sinha et al. , 2020 ) . We extend the ideas of compositional generalization to RL , and introduce the separation of zero-shot compositional generalization and fast adaptation , which is particularly relevant in RL . Compositional RL A handful of works have considered functionally compositional RL , either assuming full knowledge of the correct compositional structure of the tasks ( Devin et al. , 2017 ) or automatically learning the structure simultaneously with the modules themselves ( Goyal et al. , 2019 ; Mittal et al. , 2020 ; Yang et al. , 2020 ) . We propose a method that can handle both of these settings . Two main aspects distinguish our work from existing approaches : 1 ) our method works in a lifelong setting , where tasks arrive sequentially and the agent is not allowed to revisit previous tasks , and 2 ) we evaluate our models on tasks that are explicitly compositional at multiple hierarchical levels . A closely related problem that has long been studied in RL is that of temporally extended actions , or options , for hierarchical RL ( Sutton et al. , 1999 ; Bacon et al. , 2017 ) . Crucially , the problem we consider here differs in that the functional composition occurs at every time step , instead of the temporal chaining considered in the options literature . These two orthogonal problems model real settings in which composition could improve the RL process . Appendix A contains a deeper discussion of these connections . Other forms of hierarchical RL have learned state abstractions ( Dayan & Hinton , 1993 ; Dietterich , 2000 ; Vezhnevets et al. , 2017 ; Abel et al. , 2018 ) . While these are also related , they have mainly focused on a two-layer abstraction , where the policy passes the output of a state abstraction module as input to an action module . Instead , we consider arbitrarily many layers of abstraction that help the learning of both state and action representations . One additional form of composition in RL uses full task policies as the components , and combines policies together when solving a new task ( e.g. , by summing their value functions ; Todorov , 2009 ; Barreto et al. , 2018 ; Haarnoja et al. , 2018 ; Van Niekerk et al. , 2019 ; Nangue Tasse et al. , 2020 . Instead , we separate each policy itself into components , such that these components combine to form full policies . Batch RL RL from fixed data introduces a shift between the data distribution and the learned policy , which strains the capabilities of standard methods ( Levine et al. , 2020 ) . Recent techniques constrain the departure from the data distribution ( Fujimoto et al. , 2019b ; Laroche et al. , 2019 ; Kumar et al. , 2019 ) . This issue closely connects to lifelong RL : obtaining backward transfer requires modifying the policy for earlier tasks without extra experience . We exploit this connection by storing a portion of the data collected on each task and replaying it for batch RL after future tasks are trained . 3 THE PROBLEM OF LIFELONG FUNCTIONAL COMPOSITION IN RL . Complex RL problems can often be divided into easier subproblems . Recent years have taught us that , given unlimited experience , artificial agents can tackle such complex problems without any sense of their compositional structures ( Silver et al. , 2016 ; Gu et al. , 2017 ) . However , discovering the underlying subproblems and learning to solve those would enable learning with substantially less experience , especially when faced with numerous tasks that share common structure . While our work shares this premise with hierarchical RL efforts , those works do not consider the formulation we present here , where functional composition occurs at multiple hierarchical levels of abstraction . Formally , an RL problem is given by a Markov decision process ( MDP ) Z = 〈S , A , R , T , γ〉 , where S is the set of states , A is the set of actions , T : S×A 7→ S is the probability distributionP ( s′ | s , a ) of transitioning to state s′ upon executing action a in state s , R : S ×A 7→ R is the reward function measuring the goodness of a given state-action pair , and γ is the discount factor that reduces the importance of rewards obtained far in the future . The agent ’ s behavior is dictated by a ( possibly stochastic ) policy π : S 7→ A that selects the action to take at each state . The goal of the agent is to find the policy π∗ ∈ Π that maximizes the discounted long-term returns : E [ ∑∞i=0 γiR ( si , ai ) ] . An RL problem Z is a composition of subproblems F1 , F2 , . . . if its optimal policy π∗ can be constructed by combining solutions to those subproblems : π∗ ( s ) = m1 ◦ m2 ◦ · · · , where each mi ∈ M : Xi 7→ Yi is the solution to the corresponding Fi . We first consider these subproblems from an intuitive perspective , and later define them precisely . In RL , each subproblem could involve pure sensing , pure acting , or a combination of both . For instance , in our earlier robotic manipulation example , the task can be decomposed into recognizing objects ( sensing ) , detecting the target object and devising a plan to reach it ( combined ) , and actuating the robot to grasp the object ( acting ) . In lifelong compositional RL , the agent will be faced with a sequence of MDPs Z ( 1 ) , . . . , Z ( Tmax ) over its lifetime . We assume that all MDPs are compositions of different subsets from k shared subproblems F = { F1 , . . . , Fk } . The goal of the lifelong learner is to find the set of solutions to these subproblems as a set of modules M = { m1 , . . . , mk } , such that learning to solve a new problem reduces to finding how to combine these modules optimally . Each module can be viewed as a processing stage in a hierarchical processing pipeline , or equivalently as functions in a program , and the goal of the agent is to find the correct module to execute at each stage and the instantiation of that module ( i.e. , its parameters ) . We formalize the compositional RL problem as a graph G = ( V , E ) ( e.g. , Figure 1 ) whose nodes are the subproblem solutions along with the state and action spaces : V = F ⋃ S̆⋃ Ă , where S̆ = unique ( { S ( 1 ) , . . . , S ( Tmax ) } ) and Ă = unique ( { A ( 1 ) , . . . , A ( Tmax ) } ) are the unique state and action spaces . Each subproblem Fi corresponds to a latent representational space Yi , generated by the corresponding module mi : Xi 7→ Yi . For example , and object detection module could map an image to a segmentation mask , to be used by subsequent modules to make action decisisons . Similarly , the S ( t ) ’ s and A ( t ) ’ s can serve as representation spaces ( Yt ) . Problem Z ( t ) is then specified as a pair of state and action nodes ( S ( t ) , A ( t ) ) in the graph , and the goal is to find a path between those nodes corresponding to a policy π ( t ) ∗ that maximizesR ( t ) . More generally , the graph formalism allows for recurrent computations via walks with cycles , and parallel computations via concurrent multipaths ; an extended definition of multiwalks trivially captures both settings . In this work , we consider the path formulation , and we restrict the number of edges in the graph by organizing the modules into layers , as explained in Section 5.1 . This formalism is related to that of Chang et al . ( 2019 ) for the supervised compositional setting , but is adapted to RL . From this definition , we outline two concrete problem settings , based on the amount of information available to the agent . Both settings show how learning could benefit from compositional solutions . Zero-shot generalization with full information In some scenarios , the agent may have access to a task descriptor that encodes how the current task relates to others in terms of their compositional structures . This descriptor might be sufficient to combine modules into a solution ( i.e. , zero-shot generalization ) , provided that the agent has learned to map the descriptors into a solution structure . The descriptor could take different forms , such as a multi-hot encoding of the various components , natural language , or highlighting the target objects in the input image . Our experiments study multihot descriptors as a means to provide the compositional structure . Formally , we assume that the descriptor is given as an external input t ∈ T , and that there exists some function s : T ×M 7→ Π that can map this input and an optimal set of modules M into an optimal policy for the current task s ( t , M ) = π ( t ) ∗ . This would enable the agent to achieve compositional generalization : the ability to solve a new task entirely by reusing components learned from previous tasks . Fast adaptation with restricted information In other scenarios , the agent is not given the luxury of information about the compositional structure . This is common in RL , where the only supervisory signal is typically the reward . In this case , the agent would be incapable of zero-shot transfer . Instead , we measure generalization as its ability to learn from limited experience a task-specific function s ( t ) : M 7→ Π that combines existing modules into the optimal policy for the current task s ( t ) ( M ) = π ( t ) ∗ . Intuitively , if the agent has accumulated a useful set of modulesM , then we would expect it to be capable of quickly discovering how to combine and adapt them to solve new tasks .

Summary: The paper introduces an algorithm for lifelong reinforcement learning using functional neural composition. The algorithm first maps a new problem onto a composition of previously acquired modules; then, the agent trains/finetunes the selected module combination on the new task; and finally, the agent incorporates this newly acquired information into the existing modules. This algorithm is then evaluated on two domains, one multi-task lifelong gridworld domain and a multi-task lifelong robotics domain. The algorithm produces superior results to several other lifelong learning approaches.