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: **Update after authors' response:** After reading the other reviews, the updated manuscript and the authors response I think the authors have made several small improvements based on the reviewers' suggestions. Typically I would now be leaning towards a score of 7. However, this year's conference only allows for either a 6 or an 8. To me personally the paper is still slightly closer to 6 than 8, but to make a clear statement in favor of acceptance and facilitate the reviewers' discussion I have raised my score to an 8. I will clearly state my opinion (and hesitation to strongly endorse acceptance of the paper) during the reviewers' discussion.  **Summary**: The paper proposes a novel method for lifelong compositional RL. The latter is formally introduced in the paper as lifelong RL (learning to solve a continuous stream of tasks without re-visiting previously seen tasks) with families of tasks that have a known compositional structure. The paper proposes to use a modular neural architecture (consisting of layers of modules) to capture the compositional nature of the tasks to learn. Through an appropriate module-selection process the (topologically disjoint) modules become functionally disjoint. This is shown to facilitate forward-transfer (faster learning of novel task variations that follow the compositional structure). The paper investigates two module-selection mechanisms: a fixed mechanism where the correct sequence of modules for each task is known, and a brute-force combinatorial search over all module combination to find the highest scoring combination. Once modules are selected, the selected module-parameters are trained to increase current task-performance. This is then consolidated with previous task experience (in an off-line RL fashion) via batch-RL to avoid catastrophic forgetting. Experiments are shown on a set of 2D grid-world tasks, and a number of simulated robot arms with different sensory setups, and the method performs well against a number of baseline methods.  **Main contributions** 1) Definition of compositional families of RL tasks. As the paper correctly points out, a number of previous work on compositional architectures has used variations of off-the-shelf (lifelong) learning tasks, which are not explicitly compositional and a beneficial composition of modules is often not known intuitively (often not even the number of modules is intuitively clear). The two task-families introduced overcome these issues and are a nice contribution towards more meaningful experimentation. The gridworld task is conceptually easy and does not have very complicated perceptual problem or high-dimensional continuous action-spaces; the robotic-tasks on the other hand is closer to real-world applications (and the corresponding complexities). Significance: I think the tasks are well chosen and might become standard-tasks for lifelong compositional RL.  2) Proposal of a compositional architecture and training procedure. Given that the compositional structure of the tasks is precisely known, it is relatively straightforward to propose an exactly matching architecture. The main innovation is a training procedure to train such an architecture. Unfortunately the paper proposes two fairly straightforward solutions that need strong simplifying assumptions: one, the correct sequence of modules is known and applied accordingly, or two, brute-force search over all module combinations (which becomes exponentially costly with increasing numbers of modules, and requires training-data, and relies on disjoint sets of tasks initially until all modules have been trained at least on one task). Significance: the two module-selection mechanisms serve as important controls and lead to promising results. But they are also fairly straightforward conceptually, and rely on often unrealistically strong assumptions. What I would have liked to see for more impactful results is an attempt to train the architecture from scratch without brute-force search and disjoint tasks initially - one possible attempt would be to use RL to train a module selection policy, similarly to Chang 2019, which is used in the paper as a main inspiration for defining compositional RL tasks.  3) Definition of lifelong compositional RL and using techniques from offline RL to avoid catastrophic forgetting. Significance: these are sensible choices, and thus valid contributions, but both have been conceptually proposed before in slightly different settings (Chang 2019 defines compositional supervised problems and using offline RL to avoid catastrophic forgetting and even have backwards-transfer is also not very far-fetched).