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: This paper introduces and explores the use of *functional compositionality* in lifelong reinforcement learning (RL). In contrast to hierarchical reinforcement, which explores temporal compositionality (the chaining of options or subpolicies across time), functional compositionality involves assembling a novel, overall function (in the case of RL, the policy function) by composing together neural modules that perform multi-stage processing of the policy's inputs. Each module can be regarded as a small function that takes some abstract input X and maps it to an abstract output Y -- the initial set of modules will take the state space as the input X, while the final set will take the actions as the output Y. Modules can perform their computations in parallel, with their outputs concatenated together afterwards, or they can be chained together, with the output of one set of modules serving as the input of the next set. In this setup, the continual learning problem consists of two phases: first, the correct set of modules for the task must be selected and composed into the policy; then, the parameters of those neural network modules must be updated from the data generated by the agent. Selecting the incorrect set of modules will lead to poor transfer and catastrophic forgetting (as the parameters of the incorrect neural network modules get overwritten). This modularity can help improve robustness by reducing dependencies between modules -- this is achieved both through parallelism and through the need of downstream modules to generalise to novel combinations of upstream modules. The authors validate their approaches on a set of tasks with compositional structure, which a functionally compositional policy would be able to exploit, as it could substitute out modules according to the components of the current task. They compare their framework against alternative continual learning baselines and show that their framework is able to (i) avoid catastrophic forgetting; (ii) demonstrate better zero-shot transfer to new tasks; and (iii) even exhibit better backward transfer on training tasks.