The Surprising Effectiveness of Latent World Models for Continual Reinforcement Learning
Abstract
We study the use of model-based reinforcement learning methods, in particular, world models for continual reinforcement learning.
In continual reinforcement learning, an agent is required to solve one task and then another sequentially while retaining performance and preventing forgetting on past tasks.
World models offer a task-agnostic solution: they do not require knowledge of task changes.
World models are a straight-forward baseline for continual reinforcement learning for three main reasons.
Firstly, forgetting in the world model is prevented by persisting existing experience replay buffers across tasks, experience from previous tasks is replayed for learning the world model.
Secondly, they are sample efficient.
Thirdly and finally, they offer a task-agnostic exploration strategy through the uncertainty in the trajectories generated by the world model.
We show that world models are a simple and effective continual reinforcement learning baseline.
We study their effectiveness on Minigrid and Minihack continual reinforcement learning benchmarks and show that it outperforms state of the art task-agnostic continual reinforcement learning methods.
Introduction
[Figure]
There has been many recent successes in reinforcement learning (RL), such as in games [1], robotics [2] and in scientific applications [3, 4].
However these successes showcase methods for solving individual tasks.
Looking beyond, it is conjectured that truly scalable intelligent systems will additionally need to master many tasks in a continual manner [5, 6, 7].
The field of continual reinforcement learning (CRL) aims to develop agents which can solve many tasks, one and then another, continually while retaining performance on all previously seen tasks [8].
This paper aims to explore the possibility of using world models to solve CRL problems.
CRL is a multifaceted problem.
We showcase a method that satisfies the traditional CL desiderata: avoids catastrophic forgetting , achieves high average performance, and is scalable.
At the same time it is task-agnostic , namely, it does not require external task identification, neither during training nor deployment.
This is highly desirable from the practical point of view, and is also an uncommon capability.
The proposed method is built on top of DreamerV2 [9].
Due to this, it inherits two important features: excellent sample efficiency due to the model-based paradigm and implicit task detection via the recurrent networks architecture.
The latter enables that the method is task-agnostic [10].
The latent model allows processing image observation with a relatively low computational cost.
We further propose to use replay memory [11, 12, 13] to reduce forgetting and a task agnostic exploration method [14] to enable solving sparse reward and partially-observable tasks.
Such a method is straightforward to implement and tune.
We verify that it successfully solves a challenging suite of Minigrid and Minihack tasks.
More generally, our work shows that the model-based paradigm is a viable solution to continual reinforcement learning.
Preliminaries 2.1 Reinforcement Learning
The environments we consider are partially observable; the agent does not have access to the environment state s ∈ S .
A Partially Observable Markov Decision Process ( POMDP [15]) is the following tuple ( S , A , P, R, Ω , O , γ ) .
Here S and A are the sets of states and actions respectively, such that for s t , s t +1 ∈ S and a t ∈ A .
P ( s t +1 | s t , a t ) is the transition distribution and R ( a t , s t , s t +1 ) is the reward function.
The discount factor is γ ∈ (0 , 1) .
Ω is the set of observations, O : S × A →
P (Ω) is an observation function which defines a distribution over observations.
Actions a t are chosen using a policy π that maps observations to actions: Ω → A .
Let us assume we have access to the states s t and we are working with a finite horizon H .
Then the return from a state is
R t = (cid:80) Hi =
t γ ( i - t ) r ( s i , a i ) .
In RL the objective is to maximize the expected return J = E a i ∼ π [ R 1 | s 0 ] given an initial state s 0 .
One approach to maximizing expected return is to use a model-free approach to learn a policy π ϕ : S
→ A with a neural network parameterized by ϕ guided by an action-value function Q θ ( s t , a t ) with parameters θ .
Instead of learning a policy directly from experience we can employ model-based
RL (MBRL) and learn an intermediate model f , for instance a transition model s t +1 = f
( s t , a t ) from experience and learn our policy with additional experience generated from the model f [16].
Instead of working with the actual state s t our methods consider the observations o t and recurrent action-value functions and models to help better estimate states s t [17].
2.2 Continual Supervised Learning
Continual Learning (CL) is a setting whereby a model must master a set of tasks sequentially while maintaining performance across all previously learned tasks.
In supervised CL, the model is sequentially shown T tasks, denoted T τ for τ = 1 , . . .
, T .
Each task, T τ , is comprised of a dataset D τ = { ( x i , y i ) } N τ i =1 which a neural network is required to learn from, to perform predictions.
More generally, a task is denoted by a tuple comprised of the conditional and marginal distributions, { p τ ( y | x ) , p τ ( x ) } .
After task τ , the model will lose access to the training dataset for T τ , however its performance will be continually evaluated on all tasks T i for i ≤ τ .
For a comprehensive review of
CL scenarios see [18, 19].
2.3 Continual Reinforcement Learning
In continual RL the agent will have a budget of N interactions with each task environment T τ .
The agent is then required to learn a policy to maximize rewards in this environment, before interacting with a new environment and having to learn a new policy.
Each task is defined as a new POMDP ,
T τ = ( S τ , A τ , P τ , R τ , Ω τ , O τ , γ τ ) .
The agent is continually evaluated on all past and present tasks and so it is desirable for the agent’s policy to transfer to new tasks while not forgetting how to perform past tasks.
CRL is not a new problem setting [20], however its definition has evolved over time and some settings will differ from paper to paper, we employ the setting above which is related to previous recent work in CRL [21, 22, 13, 23, 24, 10, 7].
Related Work 3.1 Continual Supervised Learning
One approach to CL, referred to as regularization approaches regularizes a NN’s weights to ensure that optimizing for a new task finds a solution which is “close” to the previous task’s [21, 25, 26].
Working with functions can be easier than with NN weights and so task functions can be regularized to ensure that learning new function mappings are “close” across tasks [27, 28, 29].
By contrast, expansion approaches add new NN components to enable learning new tasks while preserving components for specific tasks [30, 31].
Memory approaches replay data from previous tasks when learning the current task.
This can be performed with a generative model [32].
Or samples from previous tasks ( memories ) [33, 34, 35].
3.2 Continual Reinforcement Learning
Seminal work in CRL, EWC [21] enables DQN [36] to be able to continually learn how to play different Atari games with limited forgetting.
EWC learns new Q-functions by regularizing the
L2 distance between the new task’s optimal weights and previous task’s optimal weights.
EWC requires additional supervision informing it of task changes to update its objective, select specific
Q-function head and select a task specific ϵ -greedy exploration schedule.
Progress and Compress [22] applies a regularization to policy and value function feature extractors for an actor-critic approach.
Alternatively LPG-FTW [37] learns an actor-critic that factorizes into task specific parameters and shared parameters.
Both methods require task supervision and make use of task-specific parameters and shared parameters.
Task agnostic methods like CLEAR [13] do not require task information to perform CRL.
CLEAR leverages experience replay buffers [11] to prevent forgetting: by using an actor-critic with V-trace importance sampling [38] of past experiences from the replay buffer.
Model-based RL approaches to CRL have been demonstrated where the model weights are generated from a hypernetwork which itself is conditioned a task embedding [39].
Recent work demonstrates that recurrent policies for POMDPs can obtain good overall performance on continuous control CRL benchmarks [10].
A number of previous works have studied transfer in multi-task RL settings where the goals within an environment change [40, 41, 42].
In particular by incorporating the task definition directly into the value function [41] and combining this with off-policy learning allows a CRL agent to solve multiple tasks continually, and generalize to new goals [43].
3.3 Continual Adaptation
Instead of focusing on remembering how to perform all past tasks, another line of research investigates quick adaptation to changes in environment.
This can be captured by using a latent variable and off-policy RL [44].
Alternatively, one can meta-learn a model such that it can then adapt quickly on new changes in environment [45].
All these works use small environment changes such as reward function changes or changes in gravity or mass of certain agent limbs for instance.
The environments which we consider contain different A , S and reward functions such as opening doors with keys or avoiding lava or crossing a river which is quite different in comparison.
Continual exploration strategies which use curiosity [46] can be added as an intrinsic reward in the face of non-stationary environments in infinite horizon MDPs [47].
Our proposed model uses Plan2Explore which has been shown to outperform curiosity based methods [14].
Another related area of research is open-ended learning which aims to build agents which generalized to unseen environments through a curriculum which starts off with easy tasks and then progresses to harder tasks thereby creating agents which can generalize [48, 49, 50].
World Models for Continual Reinforcement Learning
We leverage world models for learning tasks sequentially without forgetting.
We use DreamerV2 [9] which introduces a discrete stochastic and recurrent world model that is state of the art on numerous single-GPU RL benchmarks.
This is a good choice for CRL since the world model is trained by reconstructing state, action and reward trajectories from experience, we can thus leverage experience replay buffers which persist across tasks to prevent forgetting in the world model.
Additionally, we can train a policy in the imagination or in the generated trajectories of the world model, similar to generative experience replay methods in supervised CL which remember previous tasks by replaying generated data [32].
Thus, using a world model is also sample efficient.
Also, world models are taskagnostic and do not require external supervision, without signaling to the agent that it is interacting with a new task.
Additionally, by generating rollouts in the world model’s imagination the uncertainty in the world model’s predictions, more specifically the disagreement between predictions can be used as a task-agnostic exploration bonus.
To summarize, we propose using using model-based RL with recurrent world models as a viable method for CRL, see Algorithm 1 for an overview, which we instantiate using DreamerV2.
4.1 Learning the World Model
DreamerV2 learns a recurrent (latent) state-space world model (RSSM) which predicts the forward dynamics of the environment.
At each time step t the world model receives an observation o t and is required to reconstruct the observations ˆ o t conditioned on the previous actions a <t (in addition to reconstructing rewards and discounts).
The forward dynamics are modeled using an RNN, h t =
GRU ( h t − 1 , z t , a t ) [51] where h t is the hidden state z t are the discrete probabilistic latent states which condition the observation predictions p ( o t | z t , h t ) .
Trajectories are sampled from an experience replay buffer and so persisting the replay buffer across different tasks should alleviate forgetting in the world model [13].
4.2 Policy Learning inside the World Model
The policy π is learned inside the world model by using an actor-critic [52] while freezing the weights of the RSSM world model.
At each step t of the dream inside the RSSM world model a latent state z t is sampled, z t and the RNN hidden state condition the actor ˆ a t ∼ π ( · | z t , h t ) .
The reward ˆ r t + is predicted by the world model.
The policy, π is then used to obtain new trajectories in the real environment.
These trajectories are added to the experience replay buffer.
An initial observation ois used to start generating rollouts for policy learning.
This training regime ensures that the policy generalizes to previously seen environments through the world model which imagines trajectories.
4.3 Task Agnostic Exploration
The policy learns using the imaged trajectories from the RSSM world model, so world model’s predicted rewards are used as a signal for the agent’s policy and critic.
The policy is also used to gain experience inside the real environment.
So trajectories which the world model is uncertain how to predict indicate regions of the state and action space which should be prioritized by the policy when exploring the true environment.
Hence, the uncertainty in the world model’s trajectory prediction can be used as an additional intrinsic reward.
This idea underpins Plan2Explore [14] which naturally fits with DreamerV2.
The world model quantifies the uncertainty in the next latent state prediction by using a deep ensemble; multiple neural networks with independent weights.
Deep ensembles are a surprisingly robust baseline for uncertainty quantification [53] and the ensemble’s variance is used as an intrinsic reward.
The exploration neural networks in the ensemble are trained to predict the next RSSM latent features [ z t +1 , h t +1 ] .
The world model is frozen while the ensemble is trained.
The policy π observes the reward r = α i r i + α e r e , where r e is the extrinsic reward predicted by the world model, r i is the intrinsic reward, the latent disagreement between the next latent state predictions.
The coefficients α i and α e are ∈ [0 , 1] .
Hence the policy π can be trained inside the world model to seek regions in the state action space which the world model struggles to predict and hence when the policy is deployed in the environment it will seek these same regions in the state-action space to obtain new trajectories to train the RSSM world model.
The exploration strategy is significant for CRL since it is not task dependent unlike using DQN where each task needs an ϵ -greedy schedule [21, 23] or SAC [54] which needs an entropy regularizer per task [7].
Experiments
To test the performance of DreamerV2 as a CRL method we consider a set of challenging problems.
Firstly we use 3 Minigrid tasks [55].
We also consider one CRL benchmark from the CORA suite [24]: 8 Minihack tasks [56].
Code is available at https://anonymous.4open.science/r/ dv24crl-C594 .We use two primary baselines.
First, Impala which is a powerful deep RL method not designed for CRL [38].
Second, we consider CLEAR [13] which uses Impala as a base RL algorithm and leverages experience replay buffers to prevent forgetting and is task agnostic.
We evaluate our methods by measuring success rates over the course of learning, Fig. 2 and Fig. 3.
We also can use average performance, average forgetting and average forward transfer metrics [7] to assess the effectiveness of our proposed baseline.
Average Performance.
This measures how well a CRL method performs on all tasks at the end of the task sequence.
The task performance is p τ ( t ) =
[ − 1 , 1] for all τ < T .
Since we have a reward of +1 for completing the task and − 1 for being killed by a monster or falling into lava.
If each task is seen for N environment steps and we have T tasks and the τ -th task is seen over the interval of steps [( τ − 1) × N, τ × N ] .
The average performance metric for our continual learning agent is defined as:
[Equation]
where t f = N × T is the final timestep.
Forgetting.
The average forgetting is the performance difference after interacting with a task versus the performance at the end of the final task.
The average forgetting across all tasks is defined as:
[Equation]
where F τ = p [Equation] ( τ
[Equation]
N ) − p τ ( t f ) .
[Equation]
By definition the forgetting of the final T -th task is F T = 0 .
If a CRL agent has better performance at the end of the task sequence compared to after τ -th task at time-step τ × N then F τ < 0 .
Forward Transfer.
The forward transfer is the difference in task performance during continual learning compared to the single task performance.
The forward transfer is defined as:
[Equation]
FT τ > 0 means that the CRL agent achieves better performance on task τ during continual learning versus in isolation.
So this metric measures how well a CRL agent transfers knowledge from previous tasks when learning a new task.
5.1 Minigrid
Minigrid [55] is a challenging image based, partially observable and sparse reward environment.
The agent, in red, will get a reward of +1 when it gets to the green goal, Fig. 2.
The agent sees a small region of the Minigrid environment as observation o t .
We use 3 different tasks from Minigrid:
We continuously evaluate CRL agents on all tasks, see Fig.
2 for the success rates.
The results indicate that DreamerV2 is able to solve difficult exploration tasks like the DoorKey-9x9 which involves the agent having to pick up a key and then use the key to open a door before accessing the goal.
Additionally, since DreamerV2 trains its policy inside the world model it is more sample efficient than powerful baselines like CLEAR which need × 10 more environment interactions to be able to solve the easier Minigrid tasks SimpleCrossing-SN9 and LavaCrossing-9x9 , Table 1.
The addition of Plan2Explore enables DreamerV2 to solve these environments even more quickly, see Fig.
DreamerV2 does exhibit some forgetting of the DoorKey-9x9 task and this indicates that additional mechanisms to prevent forgetting might be needed.
From the metrics in Table 1 we can see that DreamerV2 has strong forward transfer.
From the learning curves for individual tasks Fig.
4 we can see that DreamerV2 struggles on independent task learning over the course of 0 .
75 M environment steps.
In contrast, when learning continually
DreamerV2 is able to solve all tasks indicating that it transfer knowledge from previous tasks.
This is not entirely surprising since the levels look similar and so the world model will already be able to reconstruct certain observations from one task to the next.
For DreamerV2 we use the model and hyperparameters from [9] with an experience replay buffer for world model learning of size 2 M.
For DreamerV2 + Plan2Explore we set the reward coefficients to α i = α e = 0 .
9 which was found by grid search of various single task Minihack environments over α i = α i = { 0 . 1 , 0 . 5 , 0 .
9 } we use the same policy for exploration and evaluation and learn world model by observation reconstruction only, rather than observation, reward and discount reconstruction.
We explore these design decisions using the Minihack benchmark in Appendix B.1.
For CLEAR we use an experience replay buffer size of 1 M.
[Figure]
P r o p o r t i o n o f s u cc e ss e s
[Figure]
h i g h
i s b e tt e r )
[Figure]
[Table]
5.2 Minihack
To test the limits of DreamerV2’s ability to remember previous tasks and its exploration mechanism we look at a longer set of tasks which are harder to solve in the from of Minihack [56].
The effect of using an additional exploration strategy is apparent as DreamerV2 + Plan2Explore is able to solve harder tasks than the baselines.
Additionally DreamerV2 + Plan2Explore is also able to achieve large forward transfer similarly with the Minigrid experiments.
Minihack is a set of diverse image based, sparse reward tasks based on the game of Nethack [57].
We test DreamerV2 performance on 8 tasks from Minihack [56].
In particular we consider the following tasks Room-Random-15x15-v0 , Room-Monster-15x15-v0 , Room-Trap-15x15-v0 ,
Room-Ultimate-15x15-v
River-Narrow-v0 ,
River-v0 ,
River-Monster-v
and
HideNSeek-v0 , which are a subset of the 12 Minihack tasks from the CORA CRL benchmark [24].
Each task is seen once and has a budget of 1 M environment interactions.
DreamerV2 is able to solve the easier first three Room environments Room-Random-15x15-v
Room-Monster-15x15-v0 Room-Trap-15x15-v0 however struggles to solve the harder later tasks and is susceptible to forgetting.
In contrast DreamerV2 + Plan2Explore is able to solve the harder
River-Narrow-v0 , River-v0 and River-Monster-v0 tasks and is also susceptible to forgetting of the initial Room tasks.
CLEAR in contrast is really stable to remembers how to solve the Room tasks with little forgetting, however is unable to solve the more difficult River tasks.
For DreamerV
[Table]
and DreamerV2 + Plan2Explore
we use the same design choices as described for the Minigrid experiments in Section 5.1.
For CLEAR we use a replay buffer size of 1 M transitions only.
The entire task sequence is 8 M steps so CLEAR is effective in preventing forgetting with a relatively small experience replay buffer.
From Table 2 we see that CLEAR has good performance, this can be explained from the learning curves in Fig.
3 where only the simplest Minihack Room tasks are solved and remembered.
CLEAR like DreamerV2 + Plan2Explore is also similarly able to effectively transfer knowledge from other tasks when learning a new task with impressive forward transfer.
On the other hand DreamerV and DreamerV2 + Plan2Explore achieve worse average performance since it is more susceptible to forgetting however they still have high forward transfer.
An instance of forward transfer can be seen by looking at the single task performance on Room-Trap-15x15-v0 Room-Monster-15x15-v0 and
River-Monster-v0 in Fig.
5 which isn’t as high as in the CRL setting.
One simple design choice we can make to alleviate forgetting further would simply be to increase the size of the experience replay buffer, this decreases forgetting to 0 .
01 ± 0 .
11 and increases the average performance to 0 .
48 ± 0 .for
a replay buffer size of 8 M, but at the same time this decreases forward transfer Fig.
Discussion and Future Works
We have explored the use of world models as a CRL baseline.
World models can be powerful CRL agents as they train the policy inside the world model and can thus be sample efficient.
World models are trained by using experience replay buffers and so we can prevent forgetting of past tasks by persisting the replay buffer from the current task to a new task and so on.
Importantly, the world model’s prediction uncertainty can be used as an additional intrinsic task agnostic reward to help exploration and solve difficult tasks in a task-agnostic fashion [14].
Previous CRL exploration strategies in the literature all require task information to be effective.
We use DreamerV2 as the world model [9] and we show that DreamerV2 is a powerful CRL method on two different difficult CRL benchmarks.
We show that world models can be a strong baseline for CRL problems compared to state of the art methods such as CLEAR [13].
DreamerV2 with Plan2Explore outperforms CLEAR on Minigrid.
It also can achieve comparable performance with CLEAR on Minihack.
On the one hand it exhibits more forgetting than CLEAR, on the other hand it can solve harder and more difficult exploration tasks than CLEAR.
Future work, will explore making world model less susceptible to forgetting.
Supplementary Material Appendix A Single Task experiments
To assess the forward transfer of DreamerV2 for CRL we need the performance of each task as a reference Eq.
Single task learning curves for Minigrid are shown in Fig.
4 and single task learning curves for all Minihack tasks in the CRL loop are shown in Fig.
Appendix B Further Experiments
A couple further experiments are introduced which are referenced in the main paper.
In Appendix B. we explore various design choices required for DreamerV2 + Plan2Explore to get the best performance for CRL.
Secondly, in Appendix B.2 we explore how increasing the size of the experience replay buffer size affects performance in the Minihack CRL benchmark.
B.1 DreamerV2 Ablation Experiments
We explore various design choices which come from the implementations of DreamerV2 [9] and
Plan2Explore [14].
The use of Plan2Explore as an intrinsic reward.
World model learning by reconstructing the observations
ˆ o t only and not the observations, rewards and discounts all together.
The use of the exploration policy at to evaluate the performance on all current and past tasks rather than having a separate exploration and evaluation policy.
[Figure]
o
[Figure]
π exp =
π
[Figure]
Avg.
Performance ( ↑ ) Avg.
Forgetting ( ↓ ) Avg.
Forward
[Figure]
( ↑ ) - - -
[Table]
09 ±
[Table]
[Table]
37 ±
[Table]
[Table]
56 ±
[Table]
[Figure]
[Table]
28 ±
[Table]
[Table]
13 ±
[Table]
[Table]
±
[Table]
[Figure]
[Table]
[Table]
[Table]
[Figure]
[Table]
[Figure]
The results are shown in Table 3.
We decided to pick the model in the final line to report the results in the main paper as they produce the good results on Minihack with relatively small standard errors.
B.2 Stability versus Plasticity: Increasing the Size of the Replay Buffer
By increasing the replay buffer size for world model learning for DreamerV2 + Plan2Explore we see that forgetting and average performance increases, however the forward transfer simultaneously decreases, Fig. 6.
This is an instance of the stability-plasticity trade-off [58] in continual learning neural network based systems.