B UILDING A S UBSPACE OF P OLICIES FOR S CALABLE C ONTINUAL L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
The ability to continuously acquire new knowledge and skills is crucial for autonomous agents.
Existing methods are typically based on either fixed-size models that struggle to learn a large number of diverse behaviors, or growing-size models that scale poorly with the number of tasks.
In this work, we aim to strike a better balance between an agent’s size and performance by designing a method that grows adaptively depending on the task sequence.
We introduce Continual Subspace of
Policies (CSP), a new approach that incrementally builds a subspace of policies for training a reinforcement learning agent on a sequence of tasks.
The subspace’s high expressivity allows CSP to perform well for many different tasks while growing sublinearly with the number of tasks.
Our method does not suffer from forgetting and displays positive transfer to new tasks.
CSP outperforms a number of popular baselines on a wide range of scenarios from two challenging domains, Brax (locomotion) and Continual World (manipulation).
Interactive visualizations of the subspace can be found at csp .
I NTRODUCTION
Developing autonomous agents that can continuously acquire new knowledge and skills is a major challenge in machine learning, with broad application in fields like robotics or dialogue systems.
In the past few years, there has been growing interest in the problem of training agents on sequences of tasks, also referred to as continual reinforcement learning (CRL, Khetarpal et al.
However, current methods either use fixed-size models that struggle to learn a large number of diverse behaviors (Hinton et al., 2006;
Rusu et al., 2016a; Li & Hoiem, 2018; Bengio & LeCun, 2007; Kaplanis et al., 2019;
Traoré et al., 2019; Kirkpatrick et al., 2017; Schwarz et al., 2018; Mallya & Lazebnik, 2018), or growing-size models that scale poorly with the number of tasks (Berseth et al., 2018; Cheung et al., 2019; Wortsman et al., 2020).
In this work, we introduce an adaptive-size model which strikes a better balance between performance and size , two crucial properties of continual learning systems (Veniat et al., 2020), thus scaling better to long task sequences.
Taking inspiration from the mode connectivity literature (Garipov et al., 2018; Gaya et al., 2021), we propose Continual Subspace of Policies ( CSP ), a new CRL approach that incrementally builds a subspace of policies (see Figure 1 for an illustration of our method).
Instead of learning a single policy, CSP maintains an entire subspace of policies defined as a convex hull in parameter space.
The vertices of this convex hull are called anchors, with each anchor representing the parameters of a policy.
This subspace captures a large number of diverse behaviors, enabling good performance on a wide range of settings.
At every stage of the CRL process, the best found policy for a previously seen task is represented as a single point in the current subspace ( i.e. unique convex combination of the anchors), which facilitates cheap storage and easy retrieval of prior solutions.
If a new task shares some similarities with previously seen ones, a good policy can often be found in the current subspace without increasing the number of parameters.
On the other hand, if a new task is very different from previously seen ones, CSP extends the current subspace by adding another anchor, and learns a new policy in the extended subspace.
In this case, the pool of candidate solutions in the subspace increases, allowing CSP to deal with more diverse tasks in the future.
The size of the subspace is increased only if this leads to performance gains larger than a given threshold, allowing users to specify the desired trade-off between performance and size ( i.e. number of parameters or memory cost).
[Figure]
8 Model Size 0.
A v e r age P
e
r
f
o r m an c e
FT-1 FT-L
PackNet EWC
PNN SAC-N
FT-N CSP
[Figure]
-Off
We evaluate our approach on 18 CRL scenarios from two different domains, locomotion in Brax and robotic manipulation in Continual World, a challenging CRL benchmark (Wołczyk et al., 2021).
We also compare CSP with a number of popular CRL baselines, including both fixed-size and growingsize methods.
As Figure 1b shows, CSP is competitive with the strongest existing methods while maintaining a smaller size .
CSP does not incur any forgetting of prior tasks and displays positive transfer to new ones.
We demonstrate that by increasing the threshold parameter, the model size can be significantly reduced without substantially hurting performance.
In addition, our qualitative analysis shows that the subspace captures diverse behaviors and even combinations of previously learned skills, allowing transfer to many new tasks without requiring additional parameters.
C ONTINUAL R EINFORCEMENT L EARNING
A continual RL problem is defined by a sequence of N tasks denoted t 1 ,...,t N .
Task i is defined by a Markov Decision Process (MDP) M i =
⟨S i , A
i , T i ,r i ,γ ⟩ with a set of states S i , a set of actions
A i , a transition function T i ∶ S i × A i → P(S i ) , and a reward function r i ∶ S i × A i
→ R .
Here, we consider that all the tasks have the same state and action space.
We also define a global policy Π ∶ [ 1 ..N ]×Z → (S → P(A)) which takes as input a task id i and a sequence of tasks Z , and outputs the policy which should be used when interacting with task i .
P(A) is a probability distribution over the action space.
We consider that each task i is associated with a training budget of interactions b i .
When the system switches to task t i + 1 it no longer has access to transitions from the M i .
We do not assume access to a replay buffer with transitions from prior tasks since this would lead to a significant increase in memory cost; instead, at each stage, we maintain a buffer with transitions from the current task, just like SAC (Haarnoja et al., 2018b).
C ONTINUAL S UBSPACE OF P OLICIES (CSP) 3. S UBSPACE OF P OLICIES
Our work builds on Gaya et al.
(2021) by leveraging a subspace of policies.
However, instead of using the subspace to train on a single environment and adapt to new ones at test time, we use it to efficiently learn tasks sequentially in the continual RL setting.
This requires designing a new approach for learning the subspace, as detailed below.
Let θ ∈ R d represent the parameters of a policy, which we denote as π ( a ∣ s,θ ) .
A subspace of policies is a simplex in R d defined by a set of anchors
Θ = { θ 1 ,...,θ k } , where any policy in this simplex can be represented as a convex combination of the anchor parameters.
Let the weight vector α ∈ R k + with ∥ α ∥ 1 = 1 denote the coefficients of the convex combination.
Therefore, each value of α uniquely defines a policy with π ( a ∣ s, [ α, Θ ])
=
π ( a ∣ s, ∑ α i θ i ) .
Throughout the paper, we sometimes refer to the weight vector α and its corresponding policy π ( a ∣ s, [ α, Θ ]) interchangeably.
3. L EARNING A LGORITHM
We propose to adaptively construct a subspace of policies for the continual RL setting in which an agent learns tasks sequentially.
We call our method Continual Subspace of Policies ( CSP ).
Our model builds a sequence of subspaces Θ 1 ,..., Θ N , one new subspace after each training task, with each subspace extending the previous one.
Note that each subspace Θ j is a collection of at most j anchors (or neural networks) i.e. ∣ Θ j ∣ ≤ j, ∀ j ∈ 1 ...N since the subspace grows sublinearly with the number of tasks (as explained below).
Hence, the number of anchors m of a subspace Θ j is not the same as the number of tasks j used to create Θ j .
The learned policies are represented as single points in these subspaces.
At each stage, CSP maintains both a set of anchors defining the current subspace, as well as the weights α corresponding to the best policies found for all prior tasks.
The best found policy for task i after training on the first j tasks, ∀ i ≤ j , is denoted as π ji ( a ∣ s ) and can be represented as a point in the subspace Θ j with a weight vector denoted α ji such that π ji ( a ∣ s )
=
π ( a ∣ s, [ α j i , Θ j ]) , where ∣ α ji ∣ = ∣ Θ
j ∣ ≤ j .
Given a set of anchors Θ j and a set of previously learned policies { α j 1 ,... ,α jj } , updating our model on the new task t j + 1 produces a new subspace Θ j + 1 and a new set of weights { α j + 1 1 ,... ,α j + 1 j + 1 } .
There are two possible cases.
One possibility is that the current subspace already contains a good policy for t j + 1 , so we just need to find the weight vector corresponding to a policy which performs well on t j + 1 .
The other possibility is that the current subspace does not contain a good policy for t j + 1 .
In this case, the algorithm produces a new subspace by adding one anchor to the previous one (see next section), and converts the previous policies to be compatible with this new subspace.
To achieve this, CSP operates in two phases:
Grow the current subspace by adding a new anchor and learning the best possible policy for task j + 1 in this subspace (where the previous j anchors are frozen).
Compare the quality of this policy with the best possible policy expressed in the previous subspace.
Based on this comparison, decide whether to extend the subspace to match the new one or prune it back to the previous one.
We now describe in detail the two phases of the learning process, namely how we grow the subspace and how we decide whether to extend or prune it (see Algorithm 1 for a pseudo-code).
3. G ROW THE S UBSPACE
Given the current subspace Θ j composed of m ≤ j anchors (with j being the number of tasks seen so far), a new subspace ˜ Θ
j + 1 is built as follows.
First a new anchor denoted θ j + 1 is added to the set of anchors such that ˜Θ j +
=
Θ j ⋃{ θ j + 1 } .
With all previous anchors frozen, we train the new anchor by sampling α values from a Dirichlet distribution parameterized by a vector with size j + 1 ,
Dir (U( j + 1 )) .
The new anchor θ j + 1 is updated by maximizing the expected return obtained by interacting with the new task j + 1 using policies defined by the sampled α ’s:
[Equation]
where R j + 1 ( τ ) is the return obtained on task j + 1 throughout trajectory τ which was generated using policy π ( a ∣ s, [ α, ˜Θ j + 1 )] .
Note that the anchor is trained such that not only one but all possible values of α tend to produce a good policy.
To do so, we sample different α per episode.
The resulting subspace thus aims at containing as many good policies as possible for task j + 1 .
3. E XTEND OR P RUNE THE S UBSPACE
To decide if the new anchor is kept, we propose to simply compare the best possible policy for task j + 1 in the new subspace with the best possible policy for the same task in the previous subspace ( i.e. without using the new anchor).
Each policy could be evaluated via Monte-Carlo (MC) estimates by doing additional rollouts in the environment and recording the average performance.
However,
// Grow the Subspace do
Sample
Set policy parameters θ α ← ∑[-j + 1 i = 
-]
1 α i θ i for K do
Collect and store a in B by a π θ α if
B
11 end 12 Use estimate: ; // Extend or Prune the Subspaceα old ← arg max ( α, 0 ) with α ∈ R m + , ∥ α ∥ 1 = 1 W ϕ ( α ) α new ← arg max α ∈ R m + 1 + , ∥ α ∥ 1 = 1 W ϕ ( α ) if W ϕ (⋅ ,α new ) > ( 1 + ϵ ) ⋅ W ϕ (⋅ ,α old ) then
Return: θ 1 ,...,θ j ,θ j + 1 ,α new ; // Extend 15 else
Return: θ 1 ,...,θ j ,α old ; // Prune 17 end this typically requires a large number ( e.g. millions) of interactions which may be impractical with a limited budget.
Thus, we propose an alternative procedure to make this evaluation sample efficient.
For each task j and corresponding subspace Θ j , our algorithm also learns a Q -function Q ( s,a,α ) which is trained to predict the expected return on task j for all possible states, actions, and all possible α ’s in the corresponding subspace.
This Q-function is slightly different from the classical ones in RL since it takes as an additional input the vector α .
This Q-function is reset for every new task.
Our algorithm is based on SAC (Haarnoja et al., 2018b), so for each task, we collect a replay buffer of interactions B which contains all states and actions seen by the agent while training on that task.
Thus, the Q -function Q ( s,a,α ) can help us directly estimate the quality W ( α ) of the policy represented by the weight vector α in the new subspace which can be computed as the average over all states and actions in the replay buffer:
[Equation]
It is thus possible to compute the value of α corresponding to the best policy in the extended subspace (denoted α new ∈ R m +
[Equation]
+
[Equation] [Equation]
as well as the value of α corresponding to the best policy in the previous subspace (denoted α old ∈ R m + , ∥ α old
[Equation]
[Equation]
In practice, α new and α old are estimated by uniformly sampling a number of α ’s in the corresponding subspace as well as a number of states and actions from the buffer.
The quality of the new subspace and the previous one can thus be evaluated by comparing W ( α new ) and W ( α old ) .
If W ( α new ) > ( 1 + ϵ ) ⋅ W ( α old ) , the subspace is extended to the new subspace ( i.e. the one after the grow phase): Θ j + 1 = ˜Θ j + 1 .
Otherwise, the subspace is pruned back to the old subspace ( i.e. the one before the growth phase): Θ j + 1 = Θ j .
Note that, if the subspace is extended, the previously learned policies have to be mapped in the new subspace such that α j +
i
[Equation]
∀ i ≤ j, α j + 1 i ∶= ( α ji , 0 ) and α j + 1 j + 1 ∶= α new .
If the subspace is not extended, then old values can be kept i.e. ∀ i ≤ j, α j + 1 i ∶= α ji and α j + 1 j + 1 ∶=
α old .
After finding the best α for the current task, the replay buffer and Q-function are reinitialized.
Hence, memory cost for training the Q-function is constant and there is no need to store data from previous tasks unlike other CRL approaches (Rolnick et al., 2019).
The policy π θ j + 1 and value function W ϕ are updated using SAC (Haarnoja et al., 2018b).
See Appendix C for more details.
3. S CALABILITY OF CSP
By having access to an large number of policies, the subspace is highly expressive so it can capture a wide range of diverse behaviors.
This enables positive transfer to many new tasks without the need for additional parameters.
As a consequence, the number of parameters increases sublinearly with the the number of tasks .
The final size of the subspace depends on the sequence of tasks, with longer and more diverse sequences requiring more anchors.
The speed of growth is controlled by the threshold ϵ which trades-off performance gains for memory efficiency (the higher the ϵ the more performance losses are tolerated to reduce memory costs).
See Figure 2a and Appendix C.1.
E XPERIMENTS 4. E NVIRONMENTS
We evaluate CSP on 18 CRL scenarios containing 35 different RL tasks, from two continuous control domains, locomotion in Brax (Freeman et al., 2021) and robotic manipulation in Continual World (CW, Wołczyk et al.
(2021)), a challenging CRL benchmark.
For CW, we run experiments on both the proposed sequence of 10 tasks (CW10), and on all 8 triplet sequences (CW3).
Each task in these sequences has a different reward function.
The goal of these experiments is to compare our approach with popular CRL methods on a well-established benchmark.
For Brax, we create new CRL scenarios based on 3 subdomains
Each scenario has 8 tasks and each task has specific dynamics .
We use a budget of 1M interactions for each task.
The goal of these experiments is to perform an in-depth study to separately evaluate capabilities specific to CRL agents such as forgetting, transfer, robustness, and compositionality (see
Appendix B.2 and Tables 5a, 5b for more information).
For each of the 4 capabilities, we create
CRL scenarios, one based on HalfCheetah and one based on Ant.
To further probe the effectiveness of our approach, we also create one CRL scenario with 4 tasks on the challenging Humanoid domain.
Here we use a budget of 2M interactions for each task.
The CRL scenarios are created following the protocol introduced in Wołczyk et al.
(2021) which proposes a systematic way of generating task sequences that test CRL agents along particular axes.
While Wołczyk et al.
(2021) focus on transfer and forgetting (see Appendix B.3), we also probe robustness ( e.g. adapting to environment perturbations such as action reversal), and compositionality ( e.g. combining two previously learned skills to solve a new task).
Each task in a sequence has different dynamics which are grounded in quasi-realistic situations such as increased or decreased gravity, friction, or limb lengths.
See Appendix B.1 for details on how these tasks were designed.
4. B ASELINES
We compare CSP with a number of popular CRL baselines such as PNN (Rusu et al., 2016b),
EWC (Kirkpatrick et al., 2017), P ACK N ET (Mallya & Lazebnik, 2018), FT-1 which finetunes a single model on the entire sequence of tasks, and FT-L2 which is like FT-1 with an additional Lregularization applied during finetuning.
We also compare with SAC-N which trains one model for each task from scratch.
While SAC-N avoids forgetting, it cannot transfer knowledge across tasks.
Finally, we compare with a method called FT-N which combines the best of both SAC-N and FT-1.
Just like SAC-N, it stores one model per task after training on it and just like FT-1, it finetunes the previous model to promote transfer.
However, FT-N and SAC-N scale poorly ( i.e. linearly) with the number of tasks in terms of both memory and compute, which makes them unfeasible for real-world applications.
Note that our method is not directly comparable with CLEAR (Rolnick et al., 2019) since we assume no access to data from prior tasks.
Storing data from all prior tasks (as
CLEAR does) is unfeasible for long task sequences due to prohibitive memory costs.
All methods use
SAC (Haarnoja et al., 2018b) as a base algorithm.
The means and standard deviations are computed over 10 seeds unless otherwise noted.
See Appendices A.1, A.2, and A.4 for more details about our protocol, baselines, and hyperparameters, respectively.
[Table]
4. A BLATIONS
We also perform a number of ablations to understand how the different components of CSP influence performance.
We first compare CSP with CSP-ORACLE which selects the best policy by sampling a large number of policies in the subspace and computing Monte-Carlo estimates of their returns.
These estimates are expected to be more accurate than the critic’s, so CSP-ORACLE can be considered an upper bound to CSP.
However, CSP-ORACLE is less efficient than CSP since it requires singificantly more interactions with the environment to find a policy for each task ( i.e. millions).
We also want to understand how much performance we lose, if any, by not adding one anchor per task.
To do this, we run an ablation with no threshold which always extends the subspace by adding one anchor for each new task.
In addition, we vary the threshold ϵ used to decide whether to extend the subspace based on how much performance is gained by doing so.
This analysis can shed more light on the trade-off between performance gain and memory cost as the threshold varies.
4. M ETRICS
Agents are evaluated across a range of commonly used CRL metrics (Wołczyk et al., 2021; Powers et al., 2021).
For brevity, following the notation in Section 2, we will use π ji ∶= Π ( i, [ t 1 ,...,t j ]) to denote the policy selected by Π for task i obtained after training on the sequence of tasks t 1 ,...,t j (in this order with j included).
Note that Π can also be defined for a single task i.e. π t i i ∶= Π
( i, [ t i ]) is the policy for task i after training the system only on
task t
i (for the corresponding budget b i ).
We also report forgetting and forward transfer (see Appendix A.3 for their formal definitions).
Average Performance is the average performance of the method across all tasks, after training on the entire sequence of tasks.
Let P
i
( π ) be the performance of the system on task i using policy π defined as P i ( π ) ∶=
E π, T i [∑ r i ( s,a )] , where r i ( s,a ) is the reward on task i when taking action a in state s .
Then, the final average performance across all tasks can be computed as P ∶= 1 N ∑ Ni = 1 P i
( π Ni ) .
Model Size is the final number of parameters of a method after training on all tasks, divided by the number of parameters of FT-1 (which is a single policy trained on all tasks, with size equal to that of an anchor).
For example, in the case of CSP, the model size is equivalent with the number of anchors defining the final subspace ∣ Θ N ∣ .
Similarly, in the case of FT-N and SAC-N, the model size is always equal to the number of tasks N.
Following Rusu et al.
(2016b)
, we cap PNN’s model size to
N since otherwise it grows quadratically with the number of tasks.
R ESULTS 5. P ERFORMANCE ON B RAX
Table 1 shows the aggregated results across all scenarios from HalfCheetah, Ant, and Humanoid.
The results are normalized using SAC-N’s performance.
See Appendix D.1 for individual results.
[Table]
CSP performs as well as or better than the strongest baselines which grow linearly with the number of tasks ( i.e. PNN, SAC-N, FT-N), while maintaining a much smaller size and thus memory cost.
In fact, CSP’s size grows sublinearly with the number of tasks.
At the same time, CSP is significantly better than the most memory-efficient baselines that have a relatively small size ( i.e. FT-1, FT-L2,
EWC, P ACK N ET ).
Naive methods like FT-1 have low memory costs and good transfer, but suffer from catastrophic forgetting.
In contrast, methods that aim to reduce forgetting such as FT-L2 or
EWC do so at the expense of transfer.
The only competitive methods in terms of performance are the ones where the number of parameters increases at least linearly with the number of tasks, such as
PNN, SAC-N, and FT-N. These methods have no forgetting because they store the models trained on each task.
SAC-N has no transfer since it trains each model from scratch, while FT-N promotes transfer as it finetunes the previous model.
However, due to their poor scalability, these methods are unfeasible for more challenging CRL scenarios with long and varied task sequences.
In contrast, CSP doesn’t suffer from forgetting since the best found policies for all prior tasks can be cheaply stored in the form of convex combinations of the subspace’s anchors ( i.e. vectors rather than model weights).
In addition, due to having access to a large number of policies ( i.e. all convex combinations of the anchors), CSP has good transfer to new tasks (that share some similarities with prior ones).
This allows CSP to achieve strong performance while its memory cost grows sublinearly with the number of tasks.
Nevertheless, CSP’s performance is not too far from that of
CSP-ORACLE, indicating that the critic can be used to find good policies in the subspace without requiring millions of interactions.
See Appendix C.2 for more details about the use of the critic.
5. P ERFORMANCE ON C ONTINUAL W ORLD
Table 2 shows results on CW10 (Wołczyk et al., 2021), a popular CRL benchmark.
The baselines used in Wołczyk et al.
(2021) use a separate linear layer for each task on top of a common network.
The authors recognize this as a limitation since it scales poorly with the number of tasks.
For a fair comparison, we implement CSP using a similar protocol where the anchors are represented by linear heads, so the number of heads grows adaptively depending on the performance threshold ϵ .
Instead of model size, we compare the final number of heads, since this now determines a method’s scalability with the number of tasks.
See Appendix A for more details about the experimental setup.
CSP performs almost as well as P ACK N ET which is the current state-of-the-art on CW10 (Wołczyk et al., 2021), and is significantly better than all other baselines.
At the same time, CSP uses about half the number of heads, thus being more memory efficient especially as the number of tasks increases.
Note that P ACK N ET suffers from a major limitation, namely that it requires prior knowledge of the total number of tasks in order to allocate resources.
If this information is not correctly specified,
P ACK N ET is likely to fail due to either not being expressive enough to handle many tasks or being too inefficient while learning only a few tasks (Wołczyk et al., 2021).
This makes P ACK N ET unfeasible for real-world applications where agents can face task sequences of varying lengths (including effectively infinite ones).
In contrast, CSP grows adaptively depending on the sequence of tasks, so it can handle both short and long sequences without any modification to the algorithm.
See
Appendix D.2 for additional results on Continual World, including CW3.
To summarize, CSP is competitive with the strongest CRL baselines, while having a lower memory cost since its size grows sublinearly with the number of tasks.
Thus, CSP maintains a good balance between model size and performance, which allows it to scale well to long task sequences.
5. A NALYSIS OF CSP
Varying the Threshold.
Figure 2a shows how performance and size vary with the threshold ϵ used to decide whether to extend the subspace or not.
Both the performance and the size are normalized with respect to CSP-LINEAR which always extends the subspace.
As expected, as ϵ increases, performance decreases, but so does the size.
Note that performance is still above 70% even as the
[Figure] [Figure]
size is cut by more than 70% , relative to CSP-LINEAR which trains a new set of parameters ( i.e. anchor) for each task.
Thus, CSP can drastically reduce memory costs without significantly hurting performance .
Practitioners can set the threshold to trade-off performance gains for memory efficiency i.e. the higher the ϵ the more performance losses are tolerated in order to reduce memory costs.
In practice, we found ϵ = 0 .
1 to offer a good trade-off between performance and model size.
Scalability.
Figure 2b shows how performance and size scale with the number of tasks in the sequence (on HalfCheetah’s compositionality scenario) for both CSP and FT-N which is our strongest baseline on this scenario.
CSP maintains both strong performance and small size ( i.e. low memory cost) even as the number of tasks increases.
In contrast, even if FT-N performs well on all these scenarios, its size grows linearly with the number of tasks, rendering it impractical for long task sequences.
Learning Efficiency.
Figure 2c shows the average performance for three different budgets ( i.e. number of interactions allowed for each task) on HalfCheetah’s robustness scenario, comparing
CSP with FT-1, FT-L2, and EWC.
These results demonstrate that CSP can learn efficiently even with a reduced budget, while still outperforming these baselines.
By keeping track of all convex combinations of previously learned behaviors, CSP enables good transfer to new tasks which in turn leads to efficient training on task sequences.
5. A NALYSIS OF THE S UBSPACE
Smoothness and Critic Accuracy.
Figure 3a shows a snapshot of a trained subspace, along with the expected reward of all policies in the subspace, for a given task.
The expected reward is computed using both Monte-Carlo (MC) rollouts, as well as our critic’s predictions using Equation 2.
As illustrated, the learned Q-function has a similar landscape with the MC reward i.e. the subspace is smooth and the critic’s estimates are accurate.
Diversity and Compositionality.
Figure 3b illustrates that the subspace contains behaviors composed of previously learned skills ( e.g. walk on the moon, walk with tiny feet, walk on the moon with tiny feet).
This allows CSP to reuse previously learned skills to find good policies for new tasks without the need for additional parameters.
The figure also shows that for a given task, the policies in the subspace cover the whole spectrum of rewards, thus emphasizing the diversity of the behaviors expressed by the subspace.
See Appendix C.5 for more analysis of the subspace.
R ELATED W ORK
Continual Reinforcement Learning.
CRL methods aim to avoid catastrophic forgetting, as well as enable transfer to new tasks, while remaining scalable to a large number of tasks.
In the past few years, multiple approaches have been proposed to address one or more of these challenges (Tessler et al., 2017; Parisi et al., 2019; Khetarpal et al., 2020; Powers et al., 2021;
Wołczyk et al., 2021;
Tessler et al., 2017; Berseth et al., 2022;
Nagabandi et al., 2018; Xie et al., 2020;
Xu et al., 2020;
Berseth et al., 2021; Ren et al., 2022; Kessler et al., 2022).
One class of methods focuses on preventing catastrophic forgetting .
Some algorithms achieve this by storing the parameters of models trained on prior tasks (Rusu et al., 2016b; Cheung et al., 2019;
Wortsman et al., 2020).
However, these methods scale poorly with the number of tasks ( i.e. at least linearly) in both compute and memory, which makes them unfeasible for more realistic scenarios.
Other methods maintain a buffer of experience from prior tasks to alleviate forgetting (Lopez-Paz & Ranzato, 2017; Isele & Cosgun, 2018; Riemer et al., 2019b; Rolnick et al., 2019; Caccia et al., 2020a).
However, this is also not scalable as the memory cost increases significantly with the number and complexity of the tasks.
In addition, many real-world applications in domains like healthcare or insurance prevent data storage due to privacy or ethical concerns.
Another class of methods focuses on improving transfer to new tasks.
Naive approaches like finetuning that train a single model on each new task provide good scalability and plasticity, but suffer from catastrophic forgetting.
To overcome this effect, methods like elastic weight consolidation (EWC) (Kirkpatrick et al., 2017) alleviate catastrophic forgetting by constraining how much the network’s weights change, but this can in turn reduce plasticity.
Another class of methods employs knowledge distillation to improve transfer in CRL (Hinton et al., 2006; Rusu et al., 2016a; Li & Hoiem, 2018; Schwarz et al., 2018;
Bengio & LeCun, 2007; Kaplanis et al., 2019; Traoré et al., 2019).
However, since these methods train a single network, they struggle to capture a large number of diverse behaviors.
There is also a large body of work which leverages the shared structure of the tasks (Xu & Zhu, 2018;
Pasunuru & Bansal, 2019; Lu et al., 2021; Mankowitz et al., 2018; Abel et al., 2017; Sodhani et al., 2021), meta-learning (Javed & White, 2019; Spigler, 2019; Beaulieu et al., 2020; Caccia et al., 2020b;
Riemer et al., 2019a;
Zhou et al., 2020; Schmidhuber, 2013), or generative models (Robins, 1995;
Silver et al., 2013; Shin et al., 2017; Atkinson et al., 2021) to improve CRL agents, but these methods don’t work very well in practice.
Lee et al.
(2020) use an iteratively expanding network but they assume no task boundaries and only consider the supervised learning setting.
Mode Connectivity and Neural Network Subspaces.
Based on prior studies of mode connectivity (Garipov et al., 2018; Draxler et al., 2018; Mirzadeh et al., 2021), Wortsman et al.
; Benton et al.
(2021) proposed the neural network subspaces to connect the solutions of different models.
More similar to our work, Doan et al.
leverages subspace properties to mitigate forgetting on a sequence of supervised learning tasks, but doesn’t focus on other aspects of the continual learning problem.
Our work was also inspired by Gaya et al.
(2021) which learns a subspace of policies for a single task for fast adaptation to new environment.
In contrast to Gaya et al. (2021), we consider the continual RL setting and adaptively grow the subspace as the agent encounters new tasks.
Our work is first to demonstrate the effectiveness of mode connectivity for the continual RL setting.
D ISCUSSION
In this paper, we propose CSP, a new continual RL method which adaptively builds a subspace of policies to learn a sequence of tasks.
CSP is competitive with the best existing methods while using significantly fewer parameters.
Thus, it strikes a good balance between performance and memory cost which allows it to scale well to long task sequences.
Our paper is first to use a subspace of policies for continual RL, and thus opens up many interesting directions for future work.
For example, one can assume a fixed size for the subspace and update all the anchors whenever the agent encounters a new task.
Another promising direction is to leverage the structure of the subspace to meta-learn or search for good policies on a given task.
Finally, active learning techniques could be employed to more efficiently evaluate policies contained in the subspace.
R EPRODUCIBILITY S TATEMENT
This work builds on the open-source RL library SaLinA (Ludovic Denoyer, 2021) and we have provided a detailed description of the experiments, environments, and our method in Appendix A, B, and C respectively.
We made the codebase available during the discussion period
A E XPERIMENTAL D ETAILS
All the experiments were implemented with SaLinA ( ? ), a flexible and simple python library for learning sequential agents.
We used Soft Actor Critic (Haarnoja et al., 2018b) as the routine algorithm for each method.
It is a popular and efficient off-policy algorithm for continuous domains.
We use a version where the entropy parameter is learned, as proposed in Haarnoja et al.
(2018a).
For each task and each methods, the twin critic networks are reset as well as the replay buffer (which has a maximum size of 1M interactions).
This is a common setting, used by several CRL benchmarks (Wołczyk et al., 2021; Powers et al., 2021)
A. E XPERIMENTAL P ROTOCOL
Brax scenarios Considering FT-N as the upper bound with which we want to compare, we decided to apply the following protocol for each short scenario.
First, we run a gridsearch on SAC hyper-parameters (see Table 3) on FT-N and select the best set in terms of final average performance (see Section 4 for the details about this metric).
Then, we freeze these hyper-parameters and performed a specific gridsearch for CSP and each baseline (see Table 4).
Each hyper-parameter set is evaluated over 10 seeds.
We believe this ensures fair comparisons, and even gives a slight advantage to FT-N compared to our method.
Note that we set the number of parallel environments to 128 and the number of update per step to 0 .
They can be seen as tasks constraints.
These values are reasonable, and FineTuning them would have increased the number of experiments by a lot.
Evaluations are done on 512 environments per experiment, each with different seeds, with no noise on the policy output.
Continual World We bypassed the preliminary SAC gridsearch and use the set of hyperparameters proposed in Wołczyk et al.
(2021).For the triplet scenarios, we performed a specific gridsearch for
CSP and each baseline (see Table 4).
For CW10, we added CSP framework only on the top layer (head), the backbone being handled by P ACK N ET .
In Wołczyk et al.
(2021), P ACK N ET uses one head per task, just like our baseline FT-N does.
In other terms, one can see it as a benchmark between FT-N and CSP on top of P ACK N ET .
Note that we did not performed a gridsearch on any hyperparameter in this case.
A. A RCHITECTURE AND B ASELINES
Architecture for Brax For the twin critics and policy, We use a network with 4 linear layers, each consisting of 256 neurons.
We use leaky ReLU activations (with α = 0 . 2 ) after every layer.
Architecture for Continual World We use the same architecture as Wołczyk et al.
(2021) for each method: we add Layer Normalization Ba et al.
after the first layer, followed by a tanh activation.
EWC (Kirkpatrick et al., 2017)
This regularization based method aims to protect parameters that are important for the previous tasks.
After each task, it uses the Fisher information matrix to approximate the importance of each weight.
We followed Wołczyk et al.
(2021) to compute the Fisher matrix and use an analytical derivation of it.
FT-L
This baseline is proposed by Kirkpatrick et al.
(2017) and re-used by Wołczyk et al.
It can be seen as a simplified version of EWC where the regularization coefficients for each parameters are equal to 1.
PNN (Rusu et al., 2016b)
This method creates a new network at the beginning of each task, as well as lateral networks that will take as inputs - for each hidden layer - the output of the networks trained on former tasks.
In this method, the number of parameters, training and inference times are growing quadratically with respect to the number of tasks, such that for our scenarios, the vanilla model size would have been around 40.
To overcome that, we followed author’s advice and reduced the number of hidden layers with respect to the final number of tasks such that it is no bigger than FT-N. Note that, as for P ACK N ET , this also requires to know this number in advance.
P ACK N ET (Mallya & Lazebnik, 2018)
the leading method in Continual World benchmark aims to learn on a new task and prune the networks in the end, and retrain the weights that have the highest amplitude.
This has two drawbacks : one has to allocate in advance a certain number of weight per task.
Without additional knowledge on the sequence, authors recommend to evenly split the number of weights per task, meaning that one has to know the final number of tasks.
In addition, the training procedure after having pruned the network does not require any interaction with the environment but is computationally expensive.
For completeness, we also report the results for the same model with hidden size doubled in Halfcheetah scenarios (see Table 12).
We named it P ACK N ETX 2 .
The slight increase of performance (7%) is not sufficient to outperform CSP on these scenarios.
A. A DDITIONAL M ETRICS
Note that for Brax, performance is always normalized using a reference cumulative reward (see
Table 15).
In Tables 12, 13, 14, 16 we also report two additional classical metrics from continual learning:
Forward Transfer measures how much a CRL system is able to transfer knowledge from task to task.
At task i , it compares the performance of the of the system trained on all previous tasks t 1 ,...,t i to the same model trained solely on task t
i .
This measure is defined as T i ∶= P
i ( π i i ) − P
i
( π t i i ) , and thus the forward transfer for all tasks is T ∶= 1 N ∑ Ni = 1 T i .
Forgetting evaluates how much a system has forgotten about task i after training on the full sequence of tasks.
It thus compares the performance of policy π ii with the performance of policy π Ni and is defined as F i ∶= P
i ( π ii ) − P
i
π Ni ) .
Similarly to the average transfer, we report the average forgetting across all tasks F ∶= 1 N ∑ Ni = 1 F i .
A. H YPERPARAMETER SELECTION
constraint of the environment, as explained in A.
[Table] [Table]
[Table]
A. C OMPUTE R ESOURCES
Each algorithm was trained using one Intel(R)
Xeon(R) CPU cores (E5-2698 v4 @
2.20GHz) and one NVIDIA V100 GPU.
Each run took between approximately 10 and 30 hours to complete for
Brax and 50 to 70 hours for CW10.
The total runtime depended on three factors: the domain, the computation time of the algorithm and the behavior of the policy.
B E NVIRONMENT D ETAILS B. D ESIGNING THE T ASKS
We used the flexibility of Brax physics engine (Freeman et al., 2021) [Table] and three of its continuous control environments Halfcheetah (obs dim: 18, action dim: 6), Ant (obs dim: 27, action dim: 8), and Humanoid (obs dim: 376, action dim: 17) to derive multiple tasks from them.
To do so, we tweaked multiple environment parameters (Table 6)) and tried to learn a policy on these new tasks.
From that pool of trials
[Table]
we kept tasks (see Table 7) that were both challenging and diversified.
[Table] [Table]
[Table] [Table]
B. D ESIGNING THE C ONTINUAL RL S CENARIOS
Inspired by Wołczyk et al.
(2021), we studied the relationship between these changes with a simple protocol: we learn a new task with a policy that has been pre-trained on a former task.
We drew forgetting and transfer tables for each pair of tasks (see Figure 5).
With this information, we designed 5 types of scenarios representing a particular challenge in continual learning.
Note that all tasks have a budget of 1 M interactions each and repeat their loop 2 times (i e. 8 M interactions in total) except for the Humanoid scenario that contains 4 tasks with 2 M interactions.
Forgetting Scenarios are designed such that a single policy tends to forget the former task when learning a new one.
∎
Halfcheetah: hugefeet → moon → carrystuff → rainfall ∎
Ant: normal → hugefeet → rainfall → moon 2.
Transfer Scenarios are designed such that a single policy has more difficulties to learn a new task after having learned the former one, rather than learning it from scratch.
∎
Halfcheetah:
carrystuff_hugegravity → moon → defectivesensor → hugefeet_rainfall ∎ Ant: nofeet_1_3 → nofeet_2_4 → nofeet_1_2 → nofeet_3_ 3.
Robustness Scenarios alternate between a normal task and a very different distraction task that disturbs the whole learning process of a single policy (we simply inverted the actions).
While this challenge looks particularly simple from a human perspective (a simple -1 vector applied on the output is fine to find an optimal policy in a continual setting), we figured out that the Fine-tuning policies struggle to recover good performances (the final average reward actually decreases).
∎
Halfcheetah: normal → inverted_actions → normal → inverted_actions ∎ Ant: normal → inverted_actions → normal → inverted_actions 4.
Compositional Scenarios present two first tasks that will be useful to learn the last one, but a very different distraction task is put at the third place to disturb this forward transfer.
The last task is indeed a combination of the two first tasks in the sense that it combines their particularities.
For example, if the
first task is "moon" and the second one is tinyfeet, the last one will combine moon’s gravity and feet morphological changes.
[Figure]
Humanoid
scenario is an additional scenario built with the challenging environment Humanoid to test our method in higher dimensions.
Here is the detailed sequence of the scenario.
[Figure]
[Table]
rr
[Table]
t
[Table]
ff
[Table]
rr
[Table]
ff
[Table]
g e g r a
[Table]
t
[Table]
d e f e
[Table]
hu
g e f ee t hu g e f ee t _ r a
[Table]
n f a ll m oo n n o r m a l r a i n f a ll t i n y f ee t t
i n y f ee t
[Table]
oo n
[Table] [Table]
[Table]
defectivesensors hugefeet hugefeet_rainfall moon normal rainfall tinyfeet tinyfeet_moon
F
[Table]
+0.
+1.
+0.
+0.
+0.
+1.
+0.
+0.
+0.
+0.
+0.
+0.
+0.
+0.
+0. +1. +0. +0.
+0. +0.
+0.
+0.
+0. +0.
+0. +0.
+1.
+0. +0.
+0.
+1.
+0. +0.
+0. +0.
+1.
+0.
+0. +0.
+1.
+0. +0.
+0. +0.
+1.
+0. +0.
+0.
+0. +0.
+0.
+0.
+0.
[Figure]
B. C ONTINUAL W ORLD T ASKS
The main benchmark of ContinualWorld is CW10, i.e. a sequence of 10 tasks designed by Wołczyk et al.
(2021) such that it has an average transfer/forgetting difficulty.
We also tested our model and our baselines on the 8 triplets (3-tasks scenarios) proposed by the authors.
The tasks originally come from Metaworld (Yu et al., 2019), and have a budget of 1M interactions each.
The scenarios were specially designed such that there is a positive forward transfer from task 1 to task 3, but task 2 is used as a distraction that interferes with these learning dynamics.
These tasks are designed with the 2.0 version of the open-source MuJoCo physics engine Todorov et al.
Here is the detailed sequence of the triplets:
push-v1 → window-close-v1 → hammer-v 2.
hammer-v1 → window-close-v1 → faucet-close-v 3.
window-close-v1 → handle-press-side-v1 → peg-unplug-side-v 4.
faucet-close-v1 → shelf-place-v1 → peg-unplug-side-v 5.
faucet-close-v1 → shelf-place-v1 → push-back-v 6.
stick-pull-v1 → peg-unplug-side-v1 → stick-pull-v 7.
stick-pull-v1 → push-back-v1 → push-wall-v 8.
push-wall-v1 → shelf-place-v1 → push-back-v
C D ETAILS ABOUT CSP C. S CALABILITY
By having access to an infinite number of policies, the subspace is highly expressive so it can capture a wide range of diverse behaviors.
This enables positive transfer to many new tasks without the need for training additional parameters.
As a consequence, the number of parameters scales sublinearly with the number of tasks.
The speed of growth is controlled by the threshold ϵ which defines how much performance we are willing to give up for decreasing the number of parameters (by the size of one policy network).
Practitioners can set the threshold to trade-off performance gains for memory efficiency ( i.e. the higher the ϵ the more performance losses are tolerated to reduce memory costs).
In practice, we noticed that setting ϵ = 0 .
1 allows good performance and a limited growth of parameters.
As the agent learns to solve more and more tasks, we expect the subspace to grow more slowly (or stop growing entirely) since it already contains many useful behaviors which transfer to new tasks.
On the other hand, if the agent encounters a task that is significantly different than previous ones and all other behaviors in the subspace, the subspace still has the flexibility to grow and incorporate entirely new skills.
The number of anchors in the final subspace is adaptive and depends on the sequence of tasks.
The longer and more diverse the task sequence, the more anchors are needed.
This property is important for real-world applications with open-ended interactions where it’s unlikely to know a priori how much capacity is required to express all useful skills.
C. T RAINING AND U SING THE C RITIC
The subspace critic W ϕ plays a central role in our method.
Compared to the vanilla SAC critic, we only add the convex combination α as an input (concatenated with the states and actions).
In this way, it is optimized not only to evaluate the future averaged return on ( s,a ) pairs of a single policy, but an infinity of policies, characterized by the convex combination α .
At the end of each task, the subspace critic has the difficult task to estimate by how much the new anchor policy θ j + 1 improves the performance of the current subspace Θ j .
To do so, one has to find the best combination of policies in the last subspace Θ j , calling it α old and the one in the current subspace Θ j + 1 , calling it α new .
In practice, we found that sampling 1024 random ( s,a ) pairs from the replay buffer at the end of the task allows to have an accurate estimation of the best policies.
This part does not require any new interaction with the environment.
However, we found that rolling out the top-k α in the new and former subspaces helps to find the best combination.
In practice, we found that setting k ∶= 8 with one rollout per combination is sufficient.
In tasks that have 1 M interactions and an episode horizon of 1000 , it requires to allocate 0 .
8% of the budget to the purpose of finding a good policy in the subspace, which does not significantly impact the training procedure.
C. S AMPLING P OLICIES FROM THE S UBSPACE
Yet, it is important to allow the critic to estimate α old .
During training, we noticed that sampling with a simple flat Dirichlet distribution (i.e. a uniform distribution over the simplex induced by the current subspace) is not enough to make the critic able to accurately estimate the performance of the last subspace (indeed, the chances of sampling a policy in the last subspace are almost surely 0.).
This is why we decided to sample in both the current and the last subspace.
The distribution we use is then a mixture of two Dirichlet (equal chances of sampling in the last subspace and in the current subspace).
We did not perform an ablation to see if balancing the mixture would increase performances.
We also tried to sample with a peaked distribution (concentration of the Dirichlet equal to the inverse of the number of the anchors) to see if it increased performances.
In some cases the new subspace is able to find good policies faster with this distribution.
It can be a good trade off between always choosing the last anchor and sampling uniformly.
C. I MPLEMENTATION
Our Pytorch implementation of CSP uses the nn.
ModuleList object to store anchor networks.
The additional computational cost compared to a single network is negligible during both training and inference as it is mentioned in Wortsman et al.
C. A NALYSIS OF THE S UBSPACE
The he best way to visualize the reward and critic value landscapes of the subspaces is when there are 3 anchors (see Figure 3).
To do saw, we draw 8192 evenly spaced points in the 3-dimensional simplex of R 3 , and average the return over 10 rollouts for the reward landscape, and 1024 pairs of ( s,a ) for the critic landscape.
We used the short version of the Compositional Scenario of HalfCheetah to display the results.
C. L IMITATIONS
CSP prevents forgetting of prior tasks, promotes transfer to new tasks, and scales sublinearly with the number of tasks.
Despite all these advantages, our method still has a number of limitations.
While the subspace grows only sublinearly with the number of tasks, this number is highly dependent on the task sequence.
In the worst case scenario, it increases linearly with the number of tasks.
On the other hand, the more similar the tasks are, the lower the size of the subspace needed to learn good policies for all tasks.
Thus, one important direction for future work is to learn a subspace of policies with a fixed number of anchors.
Instead of training an additional anchor for each new task, one could optimize a policy in the current subspace on the new task while ensuring that the best policies for prior tasks don’t change too much.
This could be formulated as a constrained optimization problem where all the anchors defining the subspace are updated for each new task, but some regions of the subspace are regularized to not change very much.
This would result in the subspace having different regions which are good for different tasks.
While the memory costs increase sublinearly with the number of tasks, the computational costs increase linearly with the number of tasks, in the current implementation of CSP.
This is because we train an additional anchor for each new task, which can be removed if it doesn’t significantly improve performance.
However, the computational costs can also be reduced if you have access to maximum reward on a task.
This is typically the case for sparse reward tasks where if the agent succeeds, it receives a reward of 1 and 0 otherwise.
In this case, there is no need to train one anchor per task.
Instead, one can simply find the best policy in the current subspace and compare its performance with the maximum reward to decide whether to train an additional anchor or not.
Hence, in this version of CSP (which is a minor modification of the current implementation) both memory and compute scale sublinearly with the number of tasks.
However, this assumption doesn’t always hold, so here we decided to implement the more general version of CSP.
In this work, we don’t specifically leverage the structure of the subspace in order to find good policies.
Hence, one promising research direction is to further improve transfer efficiency by leveraging the structure of the subspace to find good policies for new tasks.
For example, this could be done by finding the convex combination of the anchors which maximizes return on a given task.
Regularizing the geometry of the subspace to impose certain inductive biases could also be a fruitful direction for future work
[Figure]
D A DDITIONAL R ESULTS D. B RAX
[Figure]
A v e r a g e P e
r
f
o r m a n c e
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
A v e r a g e P e
r
f
o r m a n c e
[Figure] [Figure] [Figure] [Figure]
[Figure]
[Figure]
[Figure] A v e r a g e P e r f
o r m a n c e
[Figure] [Figure] [Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
D. C ONTINUAL W ORLD
In addition to the CW10 benchmark
[Table]
we also tried our method over the 8 triplets (3 tasks for each scenario) proposed in Wołczyk et al.
In these experiments
[Table]
we used CSP on the whole network.
To that end, we did not add the layer normalization described in Appendix A.2, since it is not trivial to design a subspace in these sets of parameters.
See Table 16 for the results.
[Table]
D. A BLATIONS
Ablations on the threshold parameter (see Figure 2) of CSP have been performed on the Robustness scenario of HalfCheetah.
Concerning the scalability ablation, we used the Compositional scenario of HalfCheetah and duplicated it in 3 versions (4 tasks, 8 tasks, 12 tasks).
Concerning the learning efficiency, we used the Robustness Scenario (short version, i.e. 4 tasks only) and ran it with a budget of 250k,500k and 1M interactions.
Results were averaged over 3 seeds.