Knowledge Retention in Continual Model-Based Reinforcement Learning
Authors: Haotian Fu
Yixiang Sun
Michael Littman
George Konidaris
Venue: ICML 2025 poster (2025)
Keywords: Deep Reinforcement learning
Model-based Reinforcement Learning
Continual Learning
World Models
Primary area: reinforcement_learning->deep_rl

Abstract: We propose DRAGO, a novel approach for continual model-based reinforcement learning aimed at improving the incremental development of world models across a sequence of tasks that differ in their reward functions but not the state space or dynamics. DRAGO comprises two key components: *Synthetic Experience Rehearsal*, which leverages generative models to create synthetic experiences from past tasks, allowing the agent to reinforce previously learned dynamics without storing data, and *Regaining Memories Through Exploration*, which introduces an intrinsic reward mechanism to guide the agent toward revisiting relevant states from prior tasks. Together, these components enable the agent to maintain a comprehensive and continually developing world model, facilitating more effective learning and adaptation across diverse environments. Empirical evaluations demonstrate that DRAGO is able to preserve knowledge across tasks, achieving superior performance in various continual learning scenarios.

Knowledge Retention in Continual Model-Based Reinforcement Learning
Haotian Fu" Yixiang Sun
Michael Litiman
George Konidaris
Abstract
specific to that environment, such as avoiding countertops and maneuvering around chairs. When deployed to new household, it must adapt to unfamiliar layouts while retain- ing its understanding of prior environments. Over time, as the robot encounters diverse cattings-homes with vary- ing furniture arrangements, hospitals with strict privacy constraints, or factories with evolving machinery must incrementally integrate new dynamics without forgetting earlier knowledge. This process aligns with the principles of continual learning, where agents progressively acquire new skills across tasks while preserving past experiences (Lange et al., 2022). However, real-world constraints often pro- hibit storing raw interaction data from prior tasks due to: 1. Storage limitations: Robots or embodied agents cannot indefinitely retain growing datasets (Hadsell et al., 2020). 2. Privacy regulations: Healthcare service robots handling sensitive data may be restricted from archiving task-specific interactions (Kemker et al., 201 8). 3. On-device deploy- ment: Deployed AI systems (e.g., smartphones) often rely on pre-trained models where the original training data is proprietary or privacy-sensitive.
We propose DRAGO, novel approach for contin- ual model-based reinforcement learning aimed at improving the incremental development of world models across a sequence of tasks that differ in their reward functions but not the state space or dy- namies. DRAGO comprises two key components: Synthetic Experience Rehearsal, which leverages generative models to create synthetic experiences from past tasks, allowing the agent to reinforce previously learned dynamics without storing data. and Regaining Memories Through Exploration, which introduces an intrinsic reward mechanism to guide the agent toward revisiting relevant states from prior tasks. Together, these components enable the agent to maintain comprehensive and continually developing world model, facil- itating more effective learning and adaptation across diverse environments. Empirical evalua- tions demonstrate that DRAGO is able to preserve knowledge across tasks, achieving superior per- formance in various continual learning scenarios.
In principle, continual MBRL would enable agents to learn generalizable world model supporting universal set of tasks. If data from all previous tasks are available, multi- task learning strategies (Fu et al., 2022) could solve this by leveraging shared dynamics. However, in storage- or constrycon settings, agents lack access to prior task data, rendering such approaches infeasible. In practice, we actually always have bounded storage. If we consider the hardest case, where no previous data is available: as shown in Figure 1 and our experiments, naive MBRL suf- fers catastrophic forgetting, overwriting critical dynamics (e.g., a robot forgetting kitchen layouts after learning living rooms). To address this, we need strategies that retain essen- tial knowledge without storing past data, enabling agents to build increasingly complete world models across tasks.
1. Introduction
Model-based Reinforcement Learning (MBRL) aims to en- hance decision- decision-making by developing a world model that captures the underlying dynamics of the environment. A robust world model allows an agent predict future states, plan actions, and adapt to new situations with minimal real- world trial and error. For MBRL to be effective in dynamic, real-world applications, the world model must incrementally learn and adapt, continually integrating new information as the agent encounters diverse environments and tasks.
Imagine an agent initially exploring a small, confined part of complex world, like a home-assistance robot learning to navigate kitchen. At first, the robot masters dynamics
Specifically, we propose DRAGO, novel continual MBRL approach designed to address catastrophic forgetting and incomplete world models in the absence of prior task data. DRAGO consists of two key components: Synthetic Expe- rience Rehearsal and Regaining Memories Through Explo- ration. Synthetic Experience Rehearsal uses a continually learned generative model enable the agent to simulate and learn from synthetic experiences that resemble those
Equal contribution 'Brown University. Correspon- dence Haotian Fu broueec..boon..dd.., Yixiang Sun
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Knowledge Retention in Continual Model-Baser Reinforcement Learning
Task
Task
Task
Task3
Task
Nalvo modal based
Continaal model based RL without forgetting
Figure 1. Comparison between the world model learned by naive continual MBRL and MBRL without forgetting. Each task requires the agent to move from the corner of one room to specific point the same room. Shaded areas represent the world model's coverage after finishing each task. Naively continually training MBRL (Left) tends to suffer the catastrophic forgetting problem the agent forgets almost everything about the first room after training in the second room (our experimental results support this claim). Our project identifies a method (Right) that helps world model preserve the knowledge of previous tasks even when the old data no longer available.
from prior tasks. This process allows the agent to synthesize representative transitions that resemble prior experience, reinforeing its understanding of previously learned dynam- ics without requiring access to past data. In the Regaining Memories Through Exploration component, we introduce an intrinsic reward mechanism that encourages the agent to actively explore states where the previous transition model performs well. This exploration bridges the gap between tasks by discovering connections within the environment, leading to more comprehensive and cohesive world model. To sum up, make the following contributions:
2. Background
In reinforcement learning, agent interacts with an envi- ronment modeled as a Markov Decision Process (MDP). An MDP is defined by a tuple < S, T,r,r, >, where S is the state space, A is the action space, T(s' 8, a) represents the transition dynamics, r(s, a) is the reward function, and 7 [0, 1) denotes the discount factor. Throughout this paper, T is the parametric transition model (the learned dy- namics predictor) and denotes probability or distribution.
In continual model-based reinforcement learning, the agent is presented with a sequence of tasks Ti, T ......TT.. Tn. We as- sume the agent knows when the task switches. Each task Ti is associated with its own MDP, M, (S, A, T. ra, Pi, 7), where rs a) the task-specific reward function, and ps(s) denotes the initial state distribution for task T. Importantly, all tasks share the same transition function T(s' s, a), which defines the probability of reaching state s' E S from state E S after taking action E A. this paper, we consider the case where, in each task, the agent tends to be exposed distinct aspects of the transition dynamics and different termination states.
1. We introduce DRAGO, new approach for contin- ual mosel-based reinforcement learning that addresses catastrophic forgetting while incrementally learning world model across sequential tasks without retaining any past data.
2. We propose a generative replay mechanism that synthe- sizes "old" transitions using a learned generative model alongside frozen copy of the previously trained world model.
3. We design an intrinsic reward signal that nudges the agent toward revisiting states that the old model ex- plained well-effectively "reconnecting" current expe- riences with previously learned transitions.
The objective in continual MBRL is to efficiently solve the sequence of tasks, while learning a world model To (s' s, a). parameterized by th, that captures the shared dynamics across all tasks, allowing the agent to adapt to the task- specific objectives defined by ri and p:. A challenge arises because, during training on a new task Tú. the agent in our setting only has access the replay buffer B, {(s a, s()}.
4. Through experiments on MiniGrid and DeepMind Con- trol Suite domains, show that DRAGO substan- tially improves knowledge retention compared to stan- dard continual MBRL baselines and achieves higher transfer performance, allowing faster adaptation to en- tirely new (but related) tasks. Code available at https:/ / thub. www.ooo/o///omo.oooo ccco//iiiiggnuinnn /
2

---

Knowledge Retention in Continual Mloel-BBased Reinforcement Learning
3. DRAGO
Including actions is crucial, especially in continuous action spaces where randomly sampled actions may not correspond to meaningful behaviors. Throughout the continual learning process, we also keep one copy of the old world model learned after finishing the last task (one for all previous tasks, not one for each). Then after generating the state- action pair, we feed it into this frozen old model Tata and generate synthetic next state. The synthetic data used for training the model is generated through the following steps:
The central question in this paper is: how do we aggregate the knowledge from previous tasks and learn a increasingly complete world model without forgetting, while trying to solve a sequence of tasks using MBRL? As shown in pre- vious works (Fu et al., 2022), the agent can easily learn a general world model in a multitask/meta-leamning way as long as the access to previous tasks' memories is given. Thus, a straightforward way is to figure out an approach that is able to regain the old memories that had to be discarded. We propose DRAGO, a continual MBRL approach is com- posed of two main components: dreaming and rehearsing old memories while training on new tasks ($3.1), and re- gaining memories via actively exploration ($3.2). Then we introduce the overall algorithm in $3.3.
s' Tota(3, â), (8, ~ a),
(1)
where ( ss, a; is the distribution modeled by the gener- ative model Go with parameters 0. Tou is the frozen old world model, capturing the dynamics up to previous task.
We can express the likelihood of the entire dataset, including both real data Dj for current task Ti and synthetic data DD given the parameters & and 0, as follows:
3.1. Synthetic Experience Rehearsal
To help the agent retain knowledge from previous tasks with- out direct access to past data, we introduce a method called Synthetic Experience Rehearsal. This approach enables the agent to internally generate and learn from synthetic ex- periences that resemble those from prior tasks, effectively reinforcing its understanding of the environment's dynamics and mitigating catastrophic forgetting.
(D, D 4,,0)
II p(s' a, b (L,,)))
II (
pa(š, 0)p(s' | (,);;)).
(2)
where p(s' |s, v/)b) is the likelihood of observing s' given s and a under the transition model T.. Pa((3, â; 9) is the likelihood of generating synthetic state-action pairs from the generative model Gg. This joint likelihood captures the dependencies of the synthetic data on both the generative model parameters and the frozen transition model Told-
The concept of Synthetic Experience Rehearsal draws inspi- ration from how humans and animals replay and consolidate memories during sleep (Wilson & McNaughton, 1994). Just as dreaming allows for the consolidation of memories and learning in biological systems, our method helps the agent retain and reinforce knowledge of previous dynamics by generating and learning from synthetic experiences. Imag- ine a robot that has navigated through several rooms in a building. As it progresses to new rooms, it may begin to for- get the layouts and navigation strategies of earlier ones due to limited memory capacity and the inability to revisit those rooms. By internally generating and rehearsing synthetic experiences that mimic its interactions in earlier rooms, the robot can maintain and reinforce its knowledge of how to navigate them. This internal rehearsal helps it integrate past experiences with new ones, ensuring more comprehensive understanding of the entire environment.
The posterior distribution over the transition model param- eters vé and the generative model parameters 0 is given by Bayes' theorem:
p(1b, 0 Dr, D) x p(D., D (I, (1)))
where p(b) and p(0) are the prior distributions over the parameters.
Taking the negative logarithm of the posterior (and ignoring constants independent of & and 0), we obtain the joint loss function:
Luna 0) = - log D 0. - llgg(()) - log p(0)
= -logg(()) log ((g)
Σ
log p(s' | 8, s, wi)
(s.D
Our method leverages a generative model (which also continually learned) to produce synthetic data that aids in training the dynamics model, thereby preventing forgetting of previously learned dynamics. Note that for real-world tasks, retaining the model (neural nets) usually costs much less than retaining all the training transitions, es- pecially when the task number grows larger and larger.
Loss curentt task data
(3, 0)
Σ (8,6)
log | (,);
(8,8)
Synthetic data likelihood
(4)
The dynamics model is trained by minimizing the prediction loss over the combined dataset:
Specifically, we employ generative model G that encodes and decodes both states and actions, capturing the joint dis- tribution of state-action pairs encountered in previous tasks.
Layn(w) = B(a.a,e'>>>; [[Is' - -T,,(s,;;))))) [b)
(5)
+ +XB [I a) T,(,,a
3

---

Knowledge Retention in Continual Mlodel-Based Reinforcement Learning
where À is weighting factor controlling the importance of the synthetic data loss. While this enables the agent to learn from synthetic old experience, the generative model itself (minimizing - Σ log PG( ($, å; f) in Eqn 4) also requires accumulating the knowledge of different tasks as the training goes on. Retaining such generative model for every task will also introduces huge additional cost.
within the current task. Consider the earlier example of robot exploring different rooms within building. The method introduced in the last section can generate imagined states from previously visited rooms, but without actual ex- ploration, the robot might not find the doorways or corridors connecting these rooms to its current location. Our intrinsic reward incentivizes the robot to search for these connections, enabling it to discover pathways that link the new room to the old ones. Without exploring the actual environment to find these connections, the agent's world model remains fragmented, lacking cohesive understanding of how differ- ent regions relate. To overcome this, we propose an intrinsic reward that guides the agent to:
Continual learning for the generative model. To prevent forgetting within the generative model itself, we adopt continual training strategy. We generate synthetic state- action pairs using the previous generative model G:-1;
(5 ( p(z),
and combine these with real data from the current task to form the training dataset for the new generative model: Dgen Di D. where D ((,,a)). The new gen- erative model G3 we use Variational AutoEncoder (VAE) (Kingma& & Welling, 2014)- is then trained by minimizing the loss over Dgani
Revisit Familiar States: Encourage exploration of states where the previous model T;-1 predicts accu- rately, indicating familiarity from earlier tasks.
Discover New Connections: Incentivize the agent to find paths that connect current and previous task envi- ronments, enriching the world model's completeness.
Cgen (0.) =E(ua))Doo
E Be-ge, (ala.a) po []]
(6)
KL (go, (z s, [=
Deter the agent from spending excessive time in regions where the current model T; already performs well.
Balance Learning Dynamics:
This continual learning procedure ensures that the genera- tive model retains its ability to produce state-action pairs representative of all previous tasks.
Specifically, during training on task Tr we introduce an intrinsic reward r designed to guide the agent towards states that are familiar to the previous transition model T (trained and froze after task 7;-1) but less familiar to the current model T, The intrinsic reward is defined as:
Our method is general and can be applied with other types of generative models. Additionally, integrating more sophis- ticated generative models, such as diffusion models, could further enhance the quality of synthetic experiences and improve knowledge retention in high- dineesiinnii environ- ments. We leave this for future work.
room (St, aee 9141) == a (- log 17; az) - s141l) a (- (s,,,, at) s(+1)),
3.2. Regaining Memories Through Exploration
where a denotes the sigmoid function, and a is weighting coefficient that balances the two terms.
While generating synthetic data via a generative model helps mitigate forgetting, it may not fully capture the richness of real experiences and it is subject to model error. In the meantime, to eventually build a complete world model, we would like to find a way that can "connect" knowledge gained from different tasks if they are disjoint. Thus, to further enhance the agent's retention of prior knowledge and make the world model more complete, we propose an intrinsic reward mechanism that encourages the agent to actively explore states where the previous transition model performs well, effectively "regaining" forgotten memories through real interaction with the environment, and fill in the gap between knowledge of different tasks.
Intuitively the first term assigns higher rewards when the previous transition model T;-1 predicts the next state st+1 accurately. This incentivizes the agent revisit states that were well-understood in previous tasks. The second term penalizes the agent for visiting states where the current model T; already has low prediction error. This encourages the agent to explore less familiar areas to improve the current model's understanding.
By actively exploring and connecting differrntt regions, the agent's world model becomes more comprehensive, captur- ing the dynamics across tasks more effectively. Revisiting familiar states reinforces prior knowledge, reducing the ten- dency of the model to forget previously learned information. This approach complements the synthetic data generation in Section 3.1 by providing actual experience that reinforces the agent's knowledge.
Our approach is inspired by the need to complement the generation- based rehearsal method with actual exploration that bridges the gap between different tasks. The generative model can produce states from prior tasks, but these imag- ined states might not be naturally encountered or connected
4

---

Knowledge Retention in Continual Mloel-Bassd Reinforcement Learning
Cn (-logg) To (2. -ail)
-a-o (-log (T,,, 40mmal)
7, (5.A) Po
Step
Stop Regaining Memories through Exp oraton
Symtwiic Expereecc Rehwarss!
Figure 2. The two-step process of how DRAGO retain and aggregate the knowledge Icarned from prior tasks for the world model. Step 1 involles Synthetic Experience Rehearsal, where synthetic state-action pairs are generated from the previous tasks' generative model Gz-1(=), and next states 5' are predicted using the previous transition model Ti-1. Step 2 introduces Regaining Memories through Exploration, where an intrinsic reward raom encourages the agent to explore states where the previous transition model T performs well, while penalizing states that the current model Ti already predicts accurately. Together, these components allow the agent to retain transfer knowledge across tasks.
3.3. Overall Algorithm
We implement DRAGO top of TDMPC (Hansen al., 2022) and the overall algorithm is described in Algorithm 1. Compared to regular TDMPC algorithm, we additionally train an encoder and decoder for the state-action pair as part of the generative model in $3.1. To integrate the intrinsic reward for regaining memories proposed in $3.2, we train an additional reward model, value model, and policy as a "reviewer" that aims to maximize the cumulative intrinsic reward, besides the original "learner" that aims to maximize the cumulative environmental reward. Note that the reviewer and the learner share the same world model, which is also trained using data from both.
Foomlie
From iytea Room
Roonttud
MoviB
Foorddoo
Chweshnn jemp:
jampanus
Chestal parpirs
Croutah
Walker run Walhrr waskWuker ankkbunNkkke bookwerns
Hiihher rumbback
Welker a
Wasieer buck2rur
Continual Learning
Figure 3. Visualization the evaluated domains. Task names Blue denote the continual training tasks; Task names in Red denote the test tasks. We train and test all the tasks in the order left right in the figure. E.g., we train the cheetah agent in the order of run, jump and backward. And after training on jump, we test on jump2run and jump&run.
During the inference step, DRAGO leverages Model Pre- dictive Path Integral (Williams et al., 2015) as the planning method. Given initial state and task Tú, DRAGO samples N trajectories with the world model T; and estimates the total return J, of each sampled trajectory as:
4. Experiments
We evaluated DRAGO on three continual learning domains. For each domain, we the agent train on a sequence of tasks, where the tasks share the same transition dynamics but different reward functions. Although the transition dy- namics the same, the training tasks are designed in a way such that to solve each task only part of the state space's 's transition dynamics needs to be learned and different tasks involve learning transition dynamics corresponding to dif- ferent parts of the state space with a small overlap. We evaluate the agent's continual learning performance on test tasks by measuring the agent's training performance on them, using the retained world model as initiation. The test tasks requires the combination of knowledge from more than one previously learned tasks. For example, to better transfer on Cheetah jump2run the agent is expected to still remember the knowledge learned in Cheetah run even after continual training on Cheetah jump. These transfer tasks
Ā-1
J. = E. Σ
R Ro + Q(sn
" an), 8411 Ti ag v)
where Q(-) is the learned value function. Then trajectory with the highest return is picked and the agent will execute the first action in the trajectory.
During training, the dynamics model and the generative model are trained together with the reward&value predic- tion of the learner and reviewer. At the beginning of each new task, for each new test task, randomly initial- ize the reward, policy and value models and reuse only the world model (dynamics). Moreover, unlike TDMPC, the gradients from updating Q function and reward model are detached for updating the dynamics model in DRAGO. More implementation details can be found in the appendix.
5

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
are designed to test the agent's ability retain knowledge from previous tasks, as solving them requires understanding multiple tasks.
states predictions' mean square error (MSE). The results are aligned with our intuition. Without other counter-forgeting techniques, world models easily forget almost everything learned in previous tasks and are only accurate in the transi- tion space related to the current task. By contrast, DRAGO is able to retain most of the knowledge learned in previ- ous tasks and have a increasingly complete world model as training continues, leading to the performance gain on new tasks shown in Figure 5. Note that DRAGO's performance without Synthetic Experience Rehearsal (so only has the Re- gaining Memories Through Exploration Component) drops bit compared to the full version, but it still exhibits better knowledge retention to some extent in post-task3 and post- task4, compared to naive continual TDMPC. As we also show in the ablation study, combining two components of DRAGO eventually achieves the best overall performance.
MiniGrid. We evaluated the performance of DRAGO in the MiniGrid Boiisve et al., 2023) domain using a sequence of four tasks, each set in one of the four rooms of a 27 27 gridworld. In each task, the agent starts from a fixed corner of one room, with the objective of reaching specified goal position within that room. The obstacles vary across tasks and the agent can only access other rooms by passing through a door located at the center of the grid- world, which creates a bottleneck that the agent must learn to navigate effectively in transfer tasks. Each task requires exploring a small and mostly non-overlapping portion of the world, ensuring that knowledge from one task does not directly overlap with others. To assess transfer performance, we evaluated the models learned at different stages of the continual learning process (i.e., after completing 2, 3, and 4 tasks). The evaluation was conducted on four new tasks that require the agent to move between different rooms (c.g.. start in room and move to the goal position in room 2). The tasks are designed such that solving them requires un- derstanding multiple rooms.
4.2. Overall Performance
As shown in Figure 5, we find that the proposed method DRAGO achieves the best overall performance compared to all the other approaches across three domains. The results demonstrate its advantage in continual learning settings by effectively retaining knowledge from previous tasks and transferring it to new ones. We can also see that naively continual Model-based RL may suffer from severe plasticity loss: Continual TDMPC constantly performs worse than learning from scratch baseline. Equipped with EWC, it can achieve better overall performance but still not as good as AROG But DRAGO does not fully alleviate the plasticity loss, in Cheetah Jump and runbackward (Last plot in the mid row of Figure 5), learning from scratch still has the best performance, but we can see that DRAGO still improves a lot compared to Continual TDMPC.
Deepmind Control Suite. We also evaluated the perfor- mance of DRAGO in the Cheetah and Walker domains from the Deepmind Control Suite (Tassa et al., 20. For each domain, we define a sequence of tasks that share the same dynamies but with different task goals, which requires the agent learn different parts of the state space of dynamics. Similarly, to assess transfer performance, models learned at different stages of the continual learning process. The evaluation was conducted on several new tasks that require the agent to quickly change to different locomo- tion modes from another mode (jump, run, etc.), except for two tasks: jump and runforward & jump and runbuckward, where the agent will get the maximum reward if it runs forward/backward and jumps at the same time.
evaluated the
4.3. Ablation Study
This section evaluates the essentiality of DRAGO's 's compo- nents. Specifically, we evaluate DRAGO's s performance without Synthetic Experience Rehearsal and Regaining Memories Through Exploration (reviewer) separately in four transfer tasks of Cheetah and MiniGrid. As show in Fig- ure 7, while DRAGO w/o. Rehearsal achieves similar perfor- mance with the full version in Cheetah- jangrcch-jjummannnnfofowwww the full DRAGO still has the best overall performance across domains. If we compare the performance with Contin- ual TDMPC shown in Figure 5, one single component of DRAGO consistently improves continual learning perfor- mance These results highlight the complementary roles of both components and demonstrate that each contributes significantly to mitigating forgetting and enhancing transfer capabilities in continual model-based RL settings.
We compared to baselines including: Training TDMPC from scratch for each task, continual TDMPC, where we initialize the world model with the one learned in the previ- ous task at the beginning of the new task and train with the task reward, and EWC, a regularization- based contin- ual learning method as we introduced in the related work section. We use TDMPC as the base model-based reinforce- ment learning (MBRL) algorithm for all the baselines. More experimental results can be found in appendix E & D.
4.1. Qualitative Results
In Figure 4, we also visualize the prediction accuracy of the learned world models across the whole gridworld, com- paring just naively continually training TDMPC and our method. The prediction score is calculated based on the
6

---

Knowledge Retention in Continual Model-Baser Reinforcement Learning
e30l000 betwer Donriic Phficiinn layat.
Opgled between Oyneme Treiitere and Bayen
"Surtii Eitaees Cynans Naccloraais lagett
-sannii Ietuees Syrane, Net
Continual
15
AppaSt belvent Omannic Trdictona
45160 oetrer Dowrr:
00 Hetwen Symee Predeteee Deags
DRAGO
Rehearsal
systeoeee Ouning and leyen
DRAGO
Poot 1aat-
Past-liak2
Post-laskS
Post- lack4
Figure 4. DRAGO reduces the catastrophic forgetting from previous tasks, illustrated bt the prediction score of the learned world models across the entire gridworld after each task. Light color indicates higher prediction accuracy. The heatmaps compare the performance naive continual training of TDMPC (top row). DRAGO without Synthetic Experience Rehearsal (mid row), with our proposed full DRAGO method (bottom row) after Tasks maintaining accuracy only in regions relevant to the current task, whereas DRAGO effectively retains knowledge from previous tasks.
The results show that continual MBRL suffers from significant forgetting,
MisiCral Trander BoaB afor
Mintrid Traesfes Lte 2
MiniCedd Transder 1003 after
MieiCri Transfer Ste4 ofter
Theure froa scmaad
300
20-
68
20
160093
382200
182300
150993
56099
L00KES
136803
10000
Stepe
Cheetab jump and runforward ofer Jump
Cbestub: /umpörenerwwedd ofee Jump
Chestah-jju www.:uuuwwnnnnnnnnnntoooo cuntaciounnt
Chestab Jump tal roboocward after Jump
we-
600
400-
340
208
200
100
100
2006392000ce sooce 43900cciooncie
10005 382802 400022000002 530000
199990
133000 230000 300006 4000Ce SoÓCes
Bunfewwrr22erbwerd borkward
pher
400
433-
323-
133-
433
300
203
130
18330 100KE IAKKE SM0CE SCACEE
55882 1A0A0 12000 20000 25000
(4NER 19K00 353823 20000
20301 92081 69900 09200 TEEECO
Stge
Figure evaluate continual learning transfer performance on tasks (3 domains, tasks each) that are not seen during the agent's previous training. Each plot corresponds single test task. and the agent's performance is tracked as it learns that task from scratch, using the retained world model. For each test task of MiniGrid, the agent starts in one room and have to move to the goal in another E.g., Transfer 3to4 after means that after sequentially training four tasks, the agent is tested starts room 3 and the target position is room 4. For each test task of Cheetah & Walker, the agent has to start from state in one locomotion mode and the goal switch another mode. E.g., Jump2runforward after Jump means that after training Cheetah-Jump. the agent is tested on a new task where it starts in one state of the jumping mode, and the goal is to run forward.
task where

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
MiniGrid Frundfer 1wo after
Minitrid Transfer Lw2 alwe
MinuGrid Transfer tud afler 3
Minitrad Trunsfer Btod oferr 4
DRA-0
Happle leard
10005
150000
100020
160000
50:00 100000
$2000 100020
200000
Stops
Staps
Stepa
Stops
Figure 6. Comparison of DRAGO with (1) Replay- based MBRL, which stores bounded replay buffer for past tasks (using the same memory size as DRAGO's generative model), and (2) Pseudo-rehearsal MBRL (Ketz et al., 2019), which generates old data a pretrained VAE model, without continual training of the generative model as well as our intrinsic reward mechanism. Each plot shows Minigrid transfer performance (average return vs. environment steps) on a new task after sequentially training on earlier tasks. DRAGO consistently achieves higher or faster-improving returns, suggesting stronger knowledge retention and quicker adaptation the transfer tasks.
Chentobb jerp aod rusfernset eiter 2nmp
Mariité Trasofer Lts3 alter
Maxcarr Trunsfer 2004 after
106
190
533
100
18530 Nh. Toheem
332
1000:: Se0ba: sncbon crean 320999
133m
110993
110999
32360
500.9
100900
150690
20009
Figure 7. Ablation study results on four transfer tasks in the Cheetah and MiniGrid domains, comparing the performance of DRAGO without individual components (Symtetii Experience Rehearsal and Regaining Memories Through Exploration) to the full method. While removing Rehearsal results in competitive performance in the Cheetah Cheetah-jumparnee jumpandrunforword task, the full version of DRAGO achieves superior overall performance across all tasks.
4.4. Few-shot Transfer Performance
catastrophic forgetting, where learning new tasks causes the agent's 's performance on earlier tasks to degrade due to the overwriting of important knowledge (McCloskey & Cohen, 1989). To address catastrophic forgetting, several strategies have been proposed: Regularization Based Methods (Kirk- patrick et al., 2016; Li & Hoiem, 2016; Zenke al., 2017; Nguyen al., 2017; Yu et al., 2020a): these approaches introduce constraints during training prevent significant changes to parameters important for previous tasks. Elastic Weight Consolidation (EWC) (Kirkpatrick et al., 2016) is prominent example that uses the Fisher Information Ma- trix to estimate parameter importance and penalize updates accordingly. However regularization-bare struggles in practice, especially in reinforcement learning scenarios, due to challenges in accurately estimating pa- rameter importance and scalability issues with large neural networks (Huszár, 2017; Farquhar & Gal, 2018). Replay- Based methods (Riemer et al., 2019; Rolnick et al., 2019; Oh al., 2022; Henning al., 2021; Lampinen et al., 2021): these methods typically assume unbounded storage, which is usually not feasible in practice. Our work is therefore focused the hardest case alleviating the catastrophic forgetting problem and learn a complete world model with- out prior data. Generative replay-based methods (Shin et al., 2017; Triki et al., 2017: Rao et al., 2019) share some similar high-level idea with our work. However, none of them has been applied on MBRL or World Models. In terms of Con- tinual MBRL specifically, Fu et al. (2022) show that the
We also evaluated the agent's few-shot transfer performance during the continual learning process and compared the results of DRAGO with the other baselines. The setting useful and common in real world tasks, especially for robotics, where the number of steps to interact with the envi- ronment is limited. Specifically, for each test task in Cheetah and Walker domains. we let the agent train by interacting with the environment for only 20 episodes and evaluate its average cumulative reward after training. As shown in Ta- ble 1, DRAGO outperforms the other baselines in 6 out of tasks. In the two tasks where DRAGO does not outper- form, it remains competitive, highlighting its robustness and efficiency in continual learning scenarios.
methods often
5. Related Work
Continual reinforcement learning (CRL) aims to develop agents that can learn from a sequence of tasks, retaining knowledge from previous tasks while efficiently adapting to new ones (Khetarpal et al., 2022; Abel al., 2023; Anand & Precup, 2023; Baker et al., 2023). Many recent papers investigate the plasticity loss in continual learning (Lyle et al., 2023; Abbas et al., 2023; Dohare et al., 2024). This paper focuses more on how we better retain and aggre- gate knowledge learned from previous tasks in Continual MBRL, which is related to amother central challenge in CRL,
8

---

Knowledge Retention in Continual Mlode--Based Reinforcement Learning
Average Reward
DRAGO
EWC
Continual TDMPC
Seratch
Cheetah jump2run
106.78 + 32.01
54.72 62.72
93.96 - 39.29
26.54 1: 2.67
Cheetah jump&run
248.92 = 15.38
156.98 99.68
128.58 = 100.14
182.77 + 28.58
Cheetah jump2back
331.85 = 11.05
29.93 + 7.15
73.98 = 38.45
45.15 + 4.92
Cheetah &back
147.30 = 34.29
117.92 + 1.20
140.82 = 28.00
129.75 + 20.44
Walker walk2run
332.38 1: 20.007
287.02 37.80
229.14 1 33.71
52.11 1
Walker run2back
145.98 + 17.96
150.19 2.77
128.56 9.47
60.49 9.40
Walker back2rum
229.79 + 9.77
254.09 + 70.29
241.39 + 12.64
10.76 + 18.34
Walker stand2run
265. + 8.40
177.02 + 62.48
182.71 = 30.74
64.02 : 31.54
Table 1. Comparison of few-shot transfer performance eight test tasks in Cheetah and Walker. We report the mean and standard deviation of the cumulative reward at the end of training. Bold value indicates the hest result.
agent can benefit from a joint world model for adapting to new individual tasks. Similarly, Nagabandi et al. (2019) pro- pose a meta-learning approach where dynamics model is trained to adapt quickly to tasks by learning prior over models. Hypernetwork based methods (Huang et al., 2021) have been proposed to minimize forgetting while learning task-specific parameters in the multitask setting. Liu et al. (2024) introduces pocality-sensitivy sparse encoding to learn world models incrementally in a single task online setting. Kessler al. (2023) investigate how different experience replay methods will affect the performance of MBRL. Re- lated work for model-based RL in general can be found in Appendix B.
which connects to the previous tasks), but the question of how we can better do continual learning for generative mod- els remains and leave it for future works. The current tasks tested in the paper are not highly complex, and there is limited number of tasks, which can be the reason why we do not observe this problem in our setting. Developing continual generative models can be much more challenging, but also rewarding towards the goal of real continual agent.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
6. Conclusion
We proposed DRAGO, a novel approach for continual MBRL that effectively mitigates catastrophic forgetting and enhances the transfer of knowledge across sequential tasks. By integrating Synthetic Experience Rehearsal and Regain- ing Memories Through Exploration, DRAGO retains and consolidates knowledge from previous tasks without requir- ing access to past data, resulting in a progressively more complete world model. Our empirical evaluations demon- strate that DRAGO performs well in terms of knowledge retention and transferability, making it promising solution for complex continual learning scenarios. Future work will explore extending DRAGO to larger-scale environments and more diverse task distributions.
Acknowledgement
This work supported by the Office of Naval Research (ONR) under grant #N00014-22-12592. Partial funding was also provided by The Robotics and AI Institute. The authors would like to thank David Abel, Lijing Yu, and Jingwen Zhang for their help and feedbacks on this project. This work was conducted using computational resources and services the Center for Computation and Visualization, Brown University.
References
Abbas, Z., Zhao, R., Modayil, J., White, A., and Machado,
M. C. Loss of plasticity in continual deep reinforcement learning. In Conference on Lifelong Learning Agents, 22- 25 August 2023, volume 232 of Proceedings of Machine Learning Research, pp. 220-636. PMLR, 2023.
7. Limitations
We only maintain one generative model throughout the con- tinual training process, and this could potentially have mode collapse problem as the number of the tasks grows. The generative model is expected to capture the distribution of all prior tasks, which also relies on its own generated data. Thus the forgetting issue of the generative model will ap- pear as its memory becomes "blurry" when the task number grows. To some extent, mixing the synthetic data with real world data will help mitigate this (note that the real world data can also come from the data collected by our reviewer,
Abel, D., Barreto, A., Roy, B. V., Precup, D., van Hasselt, H. P., and Singh, S. A definition of continual reinforce- ment learning. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, 2023.

---

Knowledge Retention in Continual HldellBBssd Reinforcement Learning
Anand, N. and Precup, D. Prediction and control in con- tinual reinforcement learning. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.). Advances in Neural Information Processing Sys- tems 36: Annual Conference on Neural Information Pro- cessing Systems 2023, NeurIPS 2023, 2023.
Ha, D. and Schmidhuber, J. Recurrent world models fa- cilitate policy evolution. In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Gar- nett, R. (eds.), ), Advances in Neural Information Process- ing Systems 31: Annual Conference on Neural Informa- tion Processing Systems 2018, NeurIPS 2018, pp. 2455- 2467, 2018
Baker, M. M., New, A., Aguilar-Simon, M., Al-Halah, Z., Amold, S. M. R., Ben-lwhiwhu, E., Brna, A. P., Brooks, E. A., Brown, R. C., Daniels, Z. A., Daram, A. R., De- lattre, F., Dellana, R., Faton, E., Fu, H., Grauman, K., Hostetler, J., Iqbal. S., Kent, C., Ketz, N. A., Kolouri, S., Konidaris, G. D., Kudithipudi, D., Leamed-Miller, E. G., Lee, S., Littman, M. L., Madireddy, S., Mendez, J. A. M., Nguyen, H. Q. Piatko, C. D., Pilly, P. K., Raghavan, A., Rahman, A., Ramakrishnan, S. K., Ratzlaff, N., Soltog- gio, A., Stone, P., Sur, I., Tang, Z., Tiwari, S., Vedder, K., Wang, F., Xu, Z., Yanguas-Gill A., Yedidsion, H., Yu, S., and Vallabha, G. K. domain- aggnostcc approach for characterization of lifelong learning systems. Neural networks: the official journal of the International Neural Network Society, 1 160:274 296, 2023.
Hadsell, R., Rao, D., Rusu, A. A., and Pascanu, R. Embrac- ing change:
Continual learning in deep neural networks. Trends in Cognitive Sciences, 24: 1er1100281044 2020.
Hafner, D., Lillicrap. T. P., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. Learning latent dynamics for planning from pixels. In Chaudhuri, K. and Salakhut- dinov, R. (eds.), Proceedings of the 36th International Conference on Machine Learning, ICMI.2019, volume 97 of Proceedings of Machine Learning Research, pp. 2555- 2565. PMLR, 2019.
Hafner, D., Lillicrap, T. P., Ba, J., and Norouzi, M. Dream to control: Learning behaviors by latent imagination. In ICLR. OpenReview.net, 2020.
Chevalier-Boisvettt, M., Dai, B., Towers, M., de Lazcano, R., Willems, L., Lahlou, S., Pal, S., Castro, P. S., and Terry, J. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306. 13831, 2023.
Hafner, D., Lillicrap, T. P., Norouzi, M., and Ba, J. Master- ing atari with discrete world models. In 9th International Conference on Learning Representations, ICLR 2021. OpenReview.net, 2021.
Chua, K. Calandra, R., McAllister, R., and Levine, S. Deep reinforcement learning in a handful of trials using proba- bilistic dynamics models. In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Gar- nett, R. (eds.), Advances in Neural Information Process- ing Systems 31: Annual Conference on Neural Informa- tion Processing Systems 2018, NeurlPS 2018, pp. 4759- 4770, 2018.
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T. P. Mas- tering diverse domains through world models. CoRR, abs/2301.0410 2023.
Hansen, N., Su, H.,, and Wang, X. Temporal difference learning for model predictive control. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvári, C., Niu, G., and Sabato, S. (eds.), International Conference on Machine Learning. ICMI. 2022, 17-23 July 2022, volume 162 of Proceedings of Machine Learning Research, pp. 8387-8406. PMLR, 2022.
Dohare, S., Hernandez-Garcia. J. F., Lan, Q., Rahman, P., Mahmood, A. R., and Sutton, R. S. Loss of plasticity in deep continual learning. Nat., 632((0066 2024.
Ebert, F. Finn, C., Dasari, S., Xie, A., Lee, A. X., and Levine. S. Visual foresight: Model-based deep reinforce- ment learning for vision-based robotic control. CORR, 00s/182200068,, 2018.
Hansen, N., Su, H., and Wang, X. TD-MPC2: scalable, ro- bust world models for continuous control. In The Twelfth International Conference on Learning Representations, ICLR 2024. OpenReview.net, 2024.
Farquhar, S. and Gal, Y. Towards robust evaluations of continual learning. CoRR, 005.0993,, 2018
Henning, C., Cervera, M. R., D' Angelo, F., von Oswald, J., Traber, R., Ehret, B., Kobayashi, S., Grewe, B. F,, and Sacramento, J. Posterior meta-replay for continual learning. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, pp. 14135-14149, 2021.
Fu, H., Yu, S., Littman, M., and Konidaris, G. Model- -based lifelong reinforcement learning with bayesian exploration. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Infor- mation Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, 2022.
10

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
Huang, Y., Xie, K., Bharadhwaj. H., and Shkurti, F. Con- tinual model-based reinforcement learning with hypernet- works. In IEEE International Conference on Robotics and Automation, ICRA 2021, pp. 799-805. IEEE, 2021.
Lampinen, A. K., Chan, S. C. Y., Banino, A., and Hill, F. To- wards mental time travel: hierarchical memory for rein- forcement learning agents. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, pp. 28182-281 195, 2021.
Huszár, F. Note on the quadratic penalties in elastic weight consolidation. Proceedings of the National Academy of Sciences, 115:E2496- - E2497. 2017.
Lange, M. D., Aljundi, R., Masana, M., Parisot, S., Jia, X., Leonardis, A., Slabaugh, G. G., and Tuytelaars, T. A continual learning survey: Defying forgetting in classifi- cation tasks. IEEE Trans. Pattern Anal. Mach. Intell., 44 (7):336--338) 2022.
Janner, M., Fu, J., Zhang, M., and Levine, S. When to trust your model: Model-based policy optimization. In NeurlPS, pp. 12498-12509, 12509, 2019.
Kaiser, L.. Babaeizadeh, M., Milos, P., Osinski, B., Camp- bell, R. H., Czechowski, K., Erhan, D., Finn, C., Koza- kowski, P., Levine, S., Mohiuddin, A., Sepassi, R., Tucker, G., and Michalewski, H. Model based reinforce- ment learning for atari. In ICLR. OpenReview.net, 2020.
Li, Z. and Hoiem, D. Learning without forgetting. In Leibe, B., Matas, Sehe, N., and Welling, M. (eds.), Computer Vision - ECCV 2016 14th European Conference, volume 9908 of Lecture Notes in Computer Science, pp. 614-629. Springer, 2016.
Kemker, R., McClure, M., Abitino, A., Hayes, I.,, and Kanan, C. Measuring catastrophic forgetting in neural net- works. In Mcllraith, S. A. and Weinberger, K. Q. (eds.), Proceedings of the Thirly-Second AAAI Conference Artificial Intelligence, (AAAI-18), the 30th innovative Ap- plications of Artificial Intelligence (IAAI-18), and the 8th AAAI Symposium on Educational Advances in Artifi- cial Intelligence (EAAI-18), 2018 pp. 3390-3398. AAAI Press, 2018
Liu, Z.,, Du, C., Lee, W. S., and Lin, M. Locality sensitive sparse encoding for learning world models online. In The Twelfth International Conference on Learning Represen- tations, ICLR 2024. OpenReview.net, 2024.
Lowrey, K., Rajeswaran, A., Kakade, S. M., Todorov, E., and Mordatch, I. Plan online, learn offline: Efficient learning and exploration via model-based control. In ICLR (Poster). OpenReview.net, 2019.
Kessler, S., Ostaszewski, M., Bortkiewicz, M. P., Zarski, M.. Wolczyk, M., Parker-Holden J., Roberts, S. J., and Milos, P. The effectiveness of world models for continual reinforcement learning. In Chandar, S., Pascanu, R., Sedghi, H., and Precup, D. (eds.). Conference on Lifelong Learning Agents, 22-25 August 2023, McGill University, Montréal, Québec, Canada, volume 232 of Proceedings of Machine Learning Research, pp. 184-204. PMLR, 2023.
Lyle, C., Zheng, Z., Nikishin, E., Pires, B. Á., Pascanu, R., and Dabney, W. Understanding plasticity in neural networks. In Krause, A., Brunskill, E., Cho, K., Engel- hardt, B., Sabato, S., and Scarlett, J. (eds.), International Conference Machine Learning, ICML, 2023, volume 202 of Proceedings of Machine Learning Research, pp. 13190-23211. PMLR, 2023.
McCloskey, M. and Cohen, N. J. Catastrophic interfer- ence in connectionist networks: The sequential learning problem. Psychology of Learning and Motivation, 24: 109-165, 1989.
Ketz, N., Kolouri, S., and Pilly, P. K. Using world mod- els for rseudo-rehearsal in continual learning. CoRR, abs/1903.02247, 2019.
Khetarpal, K., Riemer, M., Rish, I., and Precup, D. To- wards continual reinforcement learning: A review and perspectives. J. Artif. Intell. Res., 75:1401-1476, 2022.
Nagabandi, A., Clavera, I., Liu, S., Fearing, R. S., Abbeel, P., Levine, S., and Finn, C. Learning to adapt in dynamic, real-world environments through meta-reinforcement learning. In 7th International Conference on Learning Representations, ICLR 2019. OpenReview.net, 2019.
Kingma, D. and Welling, M. Auto-encoding variational bayes. In Bengio, Y. and LeCun, Y. (eds.), 2nd Interna- tional Conference on Learning Representations, ICLR 2014, 2014.
Nguyen, C. V., Li, Y., Bui, T. D., and Turner, R. E. Varia- tional continual learning. CoRR, abs/1710.10628. 2017.
Kirkpatrick, J., Pascanu, R., Rabinowitz, N. C., Veness,
J., Desjardins, G., Rusu, A. A., Milan, K., Quan, J., Ramalho, T., Grabska-Barwinska, A., Hassabis, D., Clopath, C., Kumaran, D., and Hadsell, R. Overcom- ing catastrophic forgetting in neural networks. CoRR, abs/1612.0079% 2016.
Nguyen, T. D., Shu, R., Pham, T., Bui, H., and Ermon, S. Temporal predictive coding for model-based planning in latent space. In ICML, volume 139 of Proceedings of Machine Learning Research, pp. 1130-8139 PMLR, 2021.

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
Oh, Y., Shin, J., Yang. E., and Hwang, S. J. Model- augmented prioritized experience replay. In The Tenth International Conference Learning Representations, ICLR 2022. OpenReview.net, 2022.
Conference on Computer Vision (ICCV), pp. 1329-1337, 2017.
Williams, G., Aldrich, A.. and Theodorou, E. A. Model
predictive path integral control using covariance variable importance sampling. CoRR, abs/1 149, 2015.
Pong. V., Gu, S., Dalal, M., and Levine, Temporal dif- ference models: Model-free deep RI. for model-based control. In ICLR (Poster). OvenReview.net, 2018.
Wilson, M. and McNaughton, B. L. Reactivation of hip- pocampal ensemble memories during sleep. Science, 265 5 5172:676-9. 9, 1994.
Rao, D., Visin, F., Rusu, A. A., Teh, Y. W., Pascanu, R., and Hadsell, Continual unsupervised representation leam- ing. In Neural Information Processing Systems, 2019.
Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., and Finn, C. Gradient surgery for multi-task learning. In Larochelle, H.,, Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances Neural Information Process- ing Systems 33: Annual Conference on Neural Informa- tion Processing Systems 2020, NeurIPS 2020, 2020a.
Riemer, M., Cases, I., Ajemian, R., Liu, M., Rish, L,, Tu, Y., and Tesauro, Learning to learn without forgetting by maximizing transfer and minimizing interference. In 7th International Conference on Learning Representations, ICLR 2019. OpenReview.net, 2019.
Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J. Y., Levine, S., Finn, C., and Ma, T. MOPO: model- offline policy optimization. In NeurIPS, 2020b.
Rolnick, D., Ahuja, A., Schwarz, J., Lillicrap, T. P., and Wayne, G. Experience replay for continual learning. In Wallach, H. M., Larochelle, H., Beygelzimer, A., d' Alché Buc, F., Fox, E. B., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information P Processing Systems 2019, NeurlPS 2019, pp. 348-358, 2019.
Zenke, F., Poole, B., and Ganguli, S. Continual learning through synaptic intelligence. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th International Con- ference on Machine Learning, ICML 2017, volume 70 of Proceedings of Machine Learning Research, pp. 3987- 3995. PMLR, 2017.
Rubinstein, R. Y. Optimization of computer simulation mod- els with rare events. European Journal of Operational Research, 99:89 1997.
Zhang. M., Vikram, S., Smith, L. M., Abbeel, P., Johnson, M. J., and Levine, S. SOLAR: deep structured latent representations for model-based reinforcement learning. CoRR, 1tt//8008.01055 2018.
Schrittwieseer, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L.,, Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., Lillicrap, T. P., and Silver, D. Mastering atari, go, chess and shogi by planning with learned model. Nat., 588(7839):604-699, 2020.
Zhang, R., Fu, H., Miao, Y., and Konidaris, G. Model-based reinforcement learning for parameterized action spaces. In Forty-first International Conference on Machine Learn- ing, ICML. 2024. OpenReview.net, 2024.
Sekar, R., Rybkin, O., Daniilidis, K., Abbeel, P., Hafner, D., and Pathak, D. Planning explore via self-supervised world models. In ICML, volume 119 of Proceedings of Machine Learning Research, pp. 9592-8592. PMLR, 2020.
Shin, H., Lee, J. K., Kim, J., and Kim, J. Continual learn- ing with deep generative replay. In Neural Information Processing Systems, 2017.
Sutton, R. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bull., 264):160-166, 1991.
Tassa, Y., Doron, Y., Muldal, A., Frez, T. Li, Y., de Las Casas, D., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., Lillicrap, T. P., and Riedmiller, M. A. Deepmind control suite. CoRR, abs/1801.0066 2018.
Triki, A., Aljundi, R., Blaschko, M. B., and Tuytelaars, T.
Encoder based lifelong learning. 20/7 IEEE International
12

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
Algorithm 1 DRAGO (Training process for each task)
Require: 30 $8, 0, o= :randomly initialized network parameters
Tota, Bold; Gola: transition network and VAE up to the previous task
11. T, X, B', B": learning ratc, coefficients, learner buffer, reviewer buffer
b. Told
fload transition model from the previous task) {load VAE encoder from the previous task) {load VAE decoder from the previous task)
2: Eg Eold
3: Go Gold
4: while not tired do
5: // Collect episode with learner and reviewer models from So
6: for step 1. = 0,..., do
7: O: ~ IT (-1sa)
(Sample with learner model]
8:
(s v2) (NV(s,, at)
{Step environment
9:
B' + B¹ Ft
{ (Add learner buffer)
10: end for
11:
for step 0 .,..., do (1+)) (s)+1, ~ (NV at)
12:
(Sample with reviewer model
13:
{Step environment
14:
rt = calculate intrinsic rewrrdss d2, S1+1)
(Equation 7] { {Add to reviewer buffer)
15:
B" + B" U (st, "e Tt: 86+1)
end for
opdate_learner and reviewwreeer( B", 0, T.
(Algorithm 2}
18:
update_vae(O,, Gola, )3)
{Algorithm 3)
19:
update_trane hawaattiinn.romsssmtttiii..c.c Tota, Gola)
(Algorithm 4}
20: end while
A. Algorithm Details
The DRAGO algorithm combines synthetic experience rehearsal and exploration-drive memory regaining to facilitate continual learning in model-based reinforcement learning (MBRL). This section provides detailed, step-by-step breakdown of DRAGO, outlining how it maintains and updates both the dynamics and generative models throughout sequence of tasks. For the first task, DRAGO exclusively trains the learner model and the rehearsal encoder-decoder pair using only online data.
1. INITIALIZATION
For each task Tú, DRAGO begins by randomly initializing the policy networks , the Q networks Q.',, and the reward networks R for both the learner and reviewer models. These components initialized separately, but they share common transition network T,
The transition network Ti, along with the synthetic experience rehearsal encoder E, and decoder G;, are initially randomly initialized for the first task. For subsequent tasks, these networks are loaded with the weights from the previous task's networks (7+-1, E5-1, and D;-1). Notably, these previously trained components (T₃-1, Eg-1, and D₃-1) are employed as fixed modules for generating synthetic data, thereby supporting the rehearsal process without further updates.
A.0.2. DATA COLLECTION
During each episode, both the learner agent and reviewer agent interact with the environment for the same number of time steps. The experiences (s, s','') encountered by each agent are stored in separate replay buffers: B for the learner and B for the reviewer. While the learner agent's 's rewards are directly sourced from the environment, the reviewer agent's intrinsic rewards are computed using the methodology outlined in Equation 7. This intrinsic reward mechanism drives the reviewer's exploration and memory regaining.
13

---

Knowledge Retention in Continual ModellBased Reinforcement Learning
Algorithm 2 update_ Jearner reeeeeeer
Require: B', B" Learner and reviewer buffers
v, " 0, o, >=: Network parameters T, A: Learning rate, coefficients
1: {s), oz ry
~ B'
(Sample trajectory from learner buffer) { (Samlle trajectory from reviewer buffer]
2: {5) 02 (7, $441 hatt ~ B"
3: o' FLtH < recculate.eviewer.ccc
4: Jg. Jo. - 0,0,0
[Initialize loss accumulation)
5: 1i = ())
6: 6 (5))).05)
7: 4 ())))
8: LQ Laluue (g)
+ Lvalue (ÿi) + Lvaaus (4))
(Calculate value loss the first observation
9: i R (61, a))
10: 9 = R (3)
11: P Rd a1)
LR Creward (
Lreward (fF) Lreward (F+)
{Calculate reward loss the first observation {Only update reward and value functions at the first step) (Initialize the estimated first observations]
13: Jg Jg + LQ LR
14: 8 = $4,8 88 i
15: for i = t....,t
t. +++ H do
16: 41 = l (8, a & = to (3)
17:
d (6), )
d. T (5,
20: Jg + JØ + X (L + (a,))
21: Jø + JΦ + X ii- "(ayyammmcs 1)) + Ldynamics (8)+
22: end for
23: 0 Φ
77 V V 24: / J V.Je
{Update online network]
{Update online network)
25: 0 (1 T)Φ T
{Update target network] {Update target network]
26: & + (1 T) TM
Algorithm 3 update. vae
Require: 0: VAE parameters
Gou: Previously trained VAE decoder
Bi: Replay buffer of the learner
1: h ~ 1)
2: (ssyath a-yuth) + Gola(h)
(Generate synthetic observations and actions}
3: hyyth + (symhh ayyith)
4: (gsyoth (synth) +- Gh
(reconstruct synthetic observation and action)
5: os s/. 1 leatH B
6: h Ea(s),ah) a1)
7: ()
Go(h)
(reconstruct sampled observation and action}
8: Jn Jo + Egen (, + Lyeal (gsynth asynth
9: 0 <- - Vojo
{update VAE parameters]
14

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
Algorithm 4 update. transition from from.syntheei synthetic data
Require: v/: Transition network parameters
Tou: Previously trained transition network
Gold: Previously trained VAE decoder
1: h 10,
2: (s a &- (h)
{Generate synthetic observations and actions) {generate next observation from old transition model]
3: s' Told (s ayyith)
4: 8 = (syrtt (
5: JØ + Jo + Laymmmiccs (8 ,')
6: Φ t v VoJe
{Update online network)
A.0.3. INFERENCE
The inference process in DRAGO is inspired by TD-MPC (Hansen et al., 2022), utilizing the Cross-Entropy Method (CEM) (Rubinstein, 1997) for action selection. During this process, a fixed number trajectories of predetermined length are sampled and simulated using the current transition model T, For each trajectory, the cumulative return is calculated. The trajectories with the highest returns, referred to as elite trajectories, are selected reshape the distribution of the initial actions. This iterative process is repeated for fixed number of iterations, ultimately yielding refined distribution over actions, which informs the final action selection. All the hyperparameters releated to the CEM algorithms is the same with TD-MPC (Hansen et al., 2022).
A.0.4. UPDATING
DRAGO updates after each episode of rollouts for the same iterations as the number of rollout time-steps The updates tries to minimize the training objective, which is the sum of several losses wegihted temporally by discount factor X. Below is a detailed description of the loss functions used in the updates:
The transition model is updated using data from both the learner agent and the reviewer agent, as well the synthetic observation-action pairs generated by the previous VAE decoder (G:-1) and the subsequent observations generated by the previous transition model (1; This process maintains the transition model's accuracy for transitions encountered in previous tasks, thereby mitigating catastrophic forgetting of the world model. Given an observation s, an action a, and a target next state s', the loss function calculates the mse between the predicted next observation using the transition model T and the next state provided:
Ldynamics c1 (s, a) s'11
However, synthetic data updates for T, only occur at fixed intervals of steps to cope with the noise arising from inaccuracies in G;-1 and 7:-1- This periodic updating strategy helps avoid noisy updates that can result from relying on outdated or inaccurate synthetic data.
Continual learning of the VAE (E, and G.) occurs concurrently with the agent's updates, Data for this learning comes from both the state-action pairs obtained from the learner model's ' rollouts and the generated sate-action pairs from G The associated loss function for the VAE Lyen is shown in Equation 6.
The reward function Rt which estimates the immediate reward from a given observation. The reward model enables the agent to estimate total return from a trajectory, and stabilizes the update for Q functions. is updated using the following loss function:
Crewand C2 Cg a4) ry
Additionally, the Q functions for both agents are updated using the TD- rD-objective shown as follows:
Evalue = c3 1 (s,, a;) - (r; + 7 (si+1,
(T
Q and R. update only using the first steps of the horizons sampled, rather than using the complete horizon in the original TD-MPC algorithm. This reduces the risk of noisy updates resulting from inaccuracies in the initial transition model.
The policy networks for both learner and reviewer agents are updated to maximize the expected Q value using
C = Qo n (x))
15

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
In the above loss functions, C1, C2 C3 are hyper parameters as weights for each losses.
A.1. Hyperparameters
Hyperparameter
Value (minigrid, cheetah, walker)
action repeat discount factor batch size maximum steps planning horizon policy fraction temperature momentum
4,
0.99
512
100, 1000, 1000
10, (25, 15), 15
0.05
0.5
0.1
reward loss coef value coef consistency loss coef vae recon loss coef vae kl loss coef temporal loss discount (p)
0.5
0.1
2
1
0.02
0.5
learning rate sampling technique
1e-3
(ER(0.6, 0.4)
target networks update freq
40, 2, 2
temperature (T)
0.01
cost coef for reviewer reward (a:)
0.5
vae latent dim
64, 256, 256
vae encoding dim
128
mlp latent dim gumble softmax temp steps per synthetic data rehearsal
512
1.0
10, 20
Table Here we list the hyperparameters used for MiniGrid World, DM-Control cheetah, and DM-Control walker. Unlisted hyperparam- eters all identical the default parameters in TD-MPC.
B. Related Work (Model-Based Reinforcement Learning)
Model-based reinforcement learning (MBRL) focuses on learning a predictive model of the environment's 's dynamics (Sutton, 1991). Learning world models (Ha & Schmidhuber, 2018; Hafner et al., 2019) specifically enables agents to accumulate knowledge about the environment's 's dynamics and generalize to new tasks or situations. By utilizing this model to simulate future states, agents can plan and make informed decisions without excessive real-world interactions. Most MBRL approaches can be categorized into two main categories in terms how the learned model is used. The first category consists of methods that the learned model generate additional data and explicitly train policy (Sutton, 1991; Pong et al., 2018; Ha & Schmidhuber, 2018; Sekar et al., 2020; Hafner et al., 2020; 2021; 2023), these approaches leverage the learned dynamics model to simulate experiences, which are then used to augment real data for policy optimization; the second category includes methods that learn the dynamics model and use it directly for planning to assign credit to actions (Ebert et al., 2018; Zhang et al., 2018; Janner et al., 2019; Hafner et al., 2019; Lowrey et al., 2019; Kaiser et al., 2020; et al., 2020b; Schrittwieser et al., 2020; Nguyen et al., 2021; Zhang et al., 2024). These methods perform online planning by simulating future trajectories using the learned model to select actions without explicitly learning a policy. Recent approaches (Hansen et al., 2022; 2024) combine both techniques and achieves superior performance various continuous control tasks. TD-MPC2 (Hansen et al., 2024) especially demonstrates the possibility of train single world model on multiple tasks at once using MBRL..
C. Tasks Specifications
Here we describe the specifications of the tasks included in this paper:
16

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
For MiniGrid World domain, all the tasks are reach goal. The pre-training tasks are dense-reward, and all fine-tuning tasks sparse-reward.
Room1to2: this task we initialize the agent inside room 1 (top left, [11, 8J) and the goal inside room 2 (top right, [14, 9D.
Roomito3: In this task we initialize the agent inside room (top left, [8, 11]) and the goal inside room 3 (bottom left,
[9, 14].
Room3to4: In this task we initialize the agent inside room 3 (bottom left, [11, 18)) and the goal inside room 4 (bottom right, [14, 17)).
For Deep Mind Control domain, all the pre-training tasks are from TD-MPC2 (Hansen et al., 2024), and the new fine-tune tasks described below:
cheetah jump2run: this task we initialize the observation as a random state when the agent is performing the task "jump", then initialize the objective be "cheetah run".
cheetah jump2hack: In this task, we initialize the observation random state when the agent is performing "jump", then initialize the objective to be "cheetah run backwards".
walker walk2run: In this task, we initialize the observation random state when the agent is performing the task "walk", then initialize the objective to be "walker run".
walker run2back: this task, we initialize the observation as a random state when the agent is performing the task "run," then initialize the objective to be "walker run backwards".
walker back2run: In this task, we initialize the observation as a random state when the agent is performing "run backwards", then initialize the objective to be "walker run".
walker stand2run: In this task. initialize the observation random state when the agent is performing the task "stand",, then initialize the objective to be "walker run".
cheetah jump& In this tasks we encourage the agent to move forward in high speed while their feet are both above the ground for longer period of time. We averaged the rewards from cheetah run and cheetah jump with a lower threshold for speed and height.
cheetah jump& back this tasks we encourage the agent to move backwards a high speed while their feet are both above the ground for a longer period of time. We averaged the rewards from cheetah run backwards and cheetah jump with a lower threshold for speed and height.
17

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
D. Additional Results of Continual Training
We investigate whether the two components we proposed have side effect on the continual training tasks, where each two of them has relatively small overlap of transition dynamics and covers different state space. As shown Table 3, DRAGO achieves similar performance with Continual TDMPC in all the training tasks, which is the MBRL baseline is built upon, demonstrating that the proposed approaches will not deteriorate the training performance or induce more plasticity loss.
Episode Reward
Cheetah run
Cheetah jump
Cheetah backward
DRAGO
652.53
587.21 646.30
624.09 580.59
Continual TDMPC
675.31
Walker run
Walker walk
Walker backward
Walker stand
708.74 693.61
954.56 959.31
953.8
972.46 982.69
956.12
Table 3. Average Episode Return of the Continual training tasks after training for 1M steps.
E. More Ablation Study Results
When trying the continual learning version for TDMPC, we find two interesting results. As shown in Figure 8 left, since we only transfer the dynamies model not the Q function, we thought excluding the Q value estimation in the planning process may yield better transfer results, but the result is the opposite. Without using the Q value in the planning process causes performance drop. Moreover, in the original TDMPC implementation, multi-step ahead prediction loss is used for updating the Q function and reward model. in the continual learning setting, we find that one-step prediction is better in complex environments like Deep Mind Control Suite as shown in the results of Cheetah-jump, which is the second one in Cheetah's 's continual training tasks.
We also investigate the influence of the frequency of synthetic experience rehearsal, the results are shown in Figure 8's second subfigure (from left right).
In Figure 8's third subfigure, we show that if we also load Cheetah run's volicy&value&rewaadt our method can reach even better results. However, this in practice requires prior knowledge that jump2run's 's reward function is similar to that of cheetah run. So it's not a scalable approach for now.
In Figure 8's last subfigure, we show a comparison of the effect of the planning horizon to the performance of DRAGO on Cheetah jump2back.
Cheeta hupmp
Chestah-Run
Chootabsjumpzzuufewwwrrr ofter Jump
Cheetel--imp2nnneee after rimbark
700
134 updata Dvg 10
coc
332
5/90
50c-
6/1
332
Costiiial TIMPC
DIACCO porcriaed wiapp
Continal TDMPC
DBAGO bedioa
plassigg withet 3
232
102
130
DRAG0 beitire
Maltitop abead prediction
08000 boriain
2.0
05
2.4
0%
0.6
1.6
1000002000131000000003
530000 060303
100003
203000
300002 400000 500003
Steps
Figure 8. More ablation study results for continual TDMPC and DRAGO
In Table 4, we compare with another baseline: Continual TDMPC + Curiosity, where we add the curiosity- Maased intrinsic reward to the continual TDMPC policy to increase exploration. We can see that DRAGO still outperforms this new continual MBRL baseline in all the four tasks. We should note that while this is reasonable baseline, the comparison is a little unfair for our method as DRAGO can also be combined with any exploration method in straightforward way. Specifically, while we have a separate reviewer model that aims to maximize our proposed intrinsic reward, our learner model that aims to solve each specific task can also be directly added with any intrinsic reward method like curiosity to encourage exploration, which does not contradict with the intrinsic reward of the separate reviewer.
We also try directly calculating the intrinsic reward of the reviewer and adding to the total reward of the learner, thus we do not need an additional reviewer model. As shown in Table 5, we see large drop of performance for the continual training
18

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
Average Reward
DRAGO
Curiusity Comiinual TDMPC
EWC
Continual TDMPC
Seratchh
Cheetah jumplrun
106.78
32.01
88.36 4 4 25.81
54.72 = 62.72
93.96
39.29 29
26.54 2.67
Choetoh jump&rrn
15.38
165.35
128.58
100.14
182.77 28.58
Chectuh junp2lark
3331.85
133.81
23.07 138.77 2 45.55
29.93
7.15
73.98
46.15 4.92
Cheetoh punpechark
147.30 34.29
117.92 = 1.20
140.82
: 28.00
129.75 20.44
Table Comparison few-shot transfer performance on four test tasks Cheetah. We report the mean and standard deviation of the cumulative reward at the end of training.
tasks, and this performance gap becomes larger and larger as the agent encounters more tasks, since it encouraged to visit more and more possibly irrelevant states. Directly adding intrinsic reward the external reward and training only one single learner model makes it hard for the agent complete the original task goal. If we only have one agent model (one policy), the intrinsic reward can have a side effect that I. discourages the agent to visit places that it is already familiar with, thus hinders it to find the optimal solution to solve the task. 2. Encourages it to visit places that the previous mode is familiar with, which could be completely irrelevant for solving the current task. By having separate reviewer policy that maximizes the intrinsic reward, we decouple the objectives. The learner policy focuses on maximizing the external reward to solve the current task effectively, while the reviewer policy explores states that help in retaining knowledge and connecting different regions of the state space. This separation allows both policies to operate without hindering each other's 's performance.
4
8
20
Figure 9. Prediction score heatmaps for replay- method (limited storage). The heatmaps illustrates the accuracy of the dynamics model during continual learning by keeping fixed part the previous replay buffer (capped similar memory size as our generative model). Compared to DRAGO's performance in Figure 4, this method does not achieve the same qualitative result.
Episode Reward
Cheetah run
Cheetah jump
Cheetah backward
DRAGO
652.53
587.24
624.09
DRAGO (Learner w. reviewer reward)
583.13
403.70
330.73
Table 5. Average Episode Return of the Continual training tasks after training for 1 IM steps.
While in all our experiments above we evaluated DRAGO using TDMPC as the MBRL baseline, we also tried to combine DRAGO with another popular model-based RL baseline PETS (Chua al., 201 2018), 8), and show the preliminary results on the same MiniGrid tasks but with dense reward (we are able make PETS work on sparse reward settings unfortunately) in Table 7. DRAGO-PETS outperforms the baseline in out of 4 tested tasks. Although the agent to perform in the union of sub observation spaces has seen in previous tasks, want to showcase the ability of ARAGO to retain knowledge by also comparing the few shot performance of pretrain tasks from only loading the world models of DRAGO and naive continual learning. Here, both world models are from a continual training of all four tasks in the Walker environment, in the order of Walker- Run, Walker- Walk, Talker-Stande and Walker- Walk- Backwards.
select finetune tasks that encourage
19

---

Knowledge Retention in Continual Model-Based Reinforcement Learning
Average Reward (Dense)
DRAGO-PETS
Continual PETS
MiniGridI after3
233.21 21.07
150.84 + 62.37
MiniGridI after2
101.03 116.21
161.70 + 44.31
MiniGridIto3 after4
138 26 99.05
43.04 + 111.93
MiniGrid3to4 after4
234.65 + 41.71
147.80 + 105.84
Table Comparison of few-shot transfer performance of PETS based methods four tasks in MiniGrid. We report the and standard deviation of the cumulative reward at the end of training.
Average Reward
DRAGO
Naive TDMPC
walker-run
371.46 + 12.79
300.52 + 27.07
walker-walk
740.89 1 54.90
685.66 1 43.05
walker-stand
871.17 55.07
463.60 1 107.12
Table Commpriiso few-shot performance of DRAGO naive TDMPC continual learning by loading world model the end pretraining on Walker tasks. We report the mean and standard deviation of the reward at 40k steps into the training.
20


Reviews:
[{"review_id": "anKImfY4jG", "reviewer": "Reviewer_5Rrm", "summary": "The authors introduce a new method (DRAGO) aimed at mitigating catastrophic forgetting in model-based RL in situations where previous experience cannot be stored. The authors propose to learn world model that compresses experience of previous tasks, and propose a novel intrinsic reward which encourages the policy to bridge the gap between different tasks, towards learning a more complete world model. The authors test their method empirically in various continual RL gridworld and continuous control domains, showing improvement over vanilla MBRL baselines and a previously successful continual-learning method (EWC).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review and for acknowledging the strong empirical results of DRAGO, as well as the clarity of our writing. We address your concerns below:\n\n\nQ: bounded memory and On-device MBRL\n\n\n\n\nA: We agree that **bounded-memory continual learning** is most critical in constrained or privacy-sensitive scenarios. In practice, real agents **cannot always store unlimited replay data**. For example, (i) *on-device learning*: robots or mobile devices have finite storage and often must learn incrementally without offloading data (for privacy or connectivity reasons), (ii) *privacy*: prior task data may contain sensitive information that cannot be archived or sent to a server. These real-world constraints motivate our setting where the agent must learn sequentially **without full past data. On-device online MBRL** addresses scenarios where a pretrained universal simulator is unavailable or too large to deploy. DRAGO’s contribution is to some extent **complementary to foundation models**: for users who have access to a foundation model, our approach could be used to continually fine-tune that model on new tasks in a memory-efficient way. Conversely, in domains not covered by a foundation model (or where data cannot leave the device), DRAGO enables continual learning from scratch. We would also like to emphasize that even though large world model pretraining is very popular right now, it is not the only correct way and not sufficient to learn a perfect world model - we would definitely need more techniques (especially continual learning) in the future and we should not stop doing research on it. \n\nQ: Evaluation Procedure\n\nA: The “test tasks” in our evaluation are **new tasks presented to the agent after the sequence of training tasks**, meant to assess how well the learned world model can be reused. We apologize for the confusion – these test tasks do involve further RL training of the policy (i.e. the agent is still learning to maximize reward on the new task), and we only load the pretrained world model’s parameters during these tests. In other words, when the agent faces a test task, only the world model is loaded to **test if it is an ideal initialization **for the new task.\n\nQ: Integration with TD-MPC, Decoder, and Learner vs. Reviewer Agents\n\nA: *Decoder* usage: DRAGO is built on TD-MPC, but we introduce a **variational generative model** (encoder–decoder) for state-action pairs as part of our Synthetic Experience Rehearsal module (Section 3.1). TD-MPC by itself uses only an encoder and latent dynamics (planning entirely in latent space), so a decoder was not needed in the original TD-MPC. In our case, however, the **decoder is essential** – it allows us to **reconstruct synthetic state-action examples** from the latent generative model of past tasks. These reconstructed experiences are fed to the world model to rehearse past dynamics. In short, **without a decoder, the agent couldn’t simulate prior states/actions explicitly**, so we added one to enable **generative replay** of past experiences (we will clarify this design choice in the text).\n\n*Learner vs. Reviewer in planning:* In Section 3.2 and Algorithm 1, we introduce two parallel actor-critic pairs – a **“learner” agent** (the original policy optimizing extrinsic reward) and a **“reviewer” agent** (an auxiliary policy optimizing an intrinsic reward)​. Both **share the same world model** and run simultaneously during training, but they have separate policy networks and reward/value heads. The learner uses the environment’s reward $r^e$ to solve the task (just like standard TD-MPC), while the reviewer receives a designed intrinsic reward $r^i$ (Equation 7 in the paper) that encourages revisiting state transitions that the **previous task’s model** could predict confidently. In practice, the two agents alternate or parallelize interactions with the environment in each training iteration (we will clarify this scheduling in the revision). During planning for action selection, each agent uses Model Predictive Control (CEM in our case) with its own reward model: the learner plans to maximize extrinsic return, while the reviewer plans to maximize intrinsic return. They do not directly interfere with each other’s action selection; they simply contribute different trajectories for training. \n\n*Why not a single combined reward/value?* We chose to keep intrinsic and extrinsic rewards separate (two policies) after initial experiments indicated that combining them can be counterproductive. A single policy optimizing a sum of extrinsic+intrinsic rewards tended to trade off one against the other, sometimes neglecting the task objective in favor of curiosity (or vice versa). By using a dedicated reviewer agent for the intrinsic objective, we ensure the extrinsic task performance remains the learner’s sole focus, while the reviewer safely explores for retention. The world model benefits from both data sources without the learner’s policy being distracted.", "claims_and_evidence": "The motivation lacks justification. For instance, are storage limitations really a bottleneck for policy learning in the present day? Can the authors qualify this with quantitative estimates? Similarly, do they have specific examples of privacy-preservation preventing the training of a generalist policy? The on-device deployment angle seems to most well justify the need for memory. However, there the justification suffers from another problem. Do the authors believe that online MBRL is likely to take place on-device, as opposed to in-context learning of some appropriately meta-trained foundation model (as is becoming standard for many AI applications)? \n\nThe results in Figures 4 and 5 are convincing that this method improves over the chosen baselines. The authors test their method on both gridworlds and continuous control tasks, and in both settings they see improvement. The ablation study in Figure 6 is also convincing. It would be useful for the authors to include information on the number of seeds they ran for each method in the figure captions. \n\nAs far as I can tell, the choice of continual RL tasks is not standard, and was rather determined by the authors. This raises the concern that these tasks could have been cherry-picked specifically to showcase the benefit of this method. Can the authors comment on why they did not use pre-defined standard benchmarks here, and whether there have any additional data or arguments that could allay a reader's concern that these tasks have been deliberately chosen to fit the specific setup in which the method is likely to succeed?", "methods_and_evaluation_criteria": "The DRAGO algorithm is well-described, and I believe that there are sufficiently many details provided that this work would be reproducible. The loss functions in equations (5) and (6) appear to be correct to me. The intrinsic reward in equation (7) is well-motivated, if a little ad-hoc. It would be useful to have a system diagram summarizing the various components of the full system (VAE, dynamics model, intrinsic reward). \n\nThe integration with TD-MPC is rather confusingly described to me. As I understand it, the philosophy of TD-MPC is to be encoder-only. However, here the authors specify that they need a decoder for state prediction. Can the authors comment on why they did not stick with the encoder-only philosophy of TD-MPC. It is also unclear to me how the \"learner\" and \"reviewer\" portions of the agent are combined when using MPC for planning. Can the authors comment on why they did not simply learn one value function on the sum of the intrinsic and extrinsic reward, and how the reviewer and learner combine to produce the policy?\n\nThis sentence is also confusing to me: For each new test task, we randomly initialize the reward, policy and value models and reuse only\nthe world model (dynamics)\". Why are the reward, policy and value models thrown away for test tasks if these are kept throughout training? Is this because of an empirical performance gain or because it would be unfair to keep them around for test, for some reason. More justification is needed here.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The evaluation setup is insufficiently well described. Is RL taking place during the \"test\" tasks or are these purely in-context learning? I assume that RL is still taking place, but then the designation of these as \"test\" tasks seems rather odd (usually one things of a train-test split, where there is no in-weights learning taking place on the test split). More explanation is required to convince me that the evaluation procedure is reasonable, rigorous and reflects a likely real-world deployment setup as motivated in the introduction.", "relation_to_broader_scientific_literature": "Both the MBRL and continual RL literatures are well-reviewed. One line of work that it may be useful to additionally mention in this regard is world models learned from data (e.g. Genie https://arxiv.org/pdf/2402.15391, UniSim https://arxiv.org/abs/2310.06114). Much as in the same way that large language models meta-learning on internet scale data has led to foundational models capable of long in-context continual adaptation, the same may be true of world models. It would be interesting to hear the authors thoughts on this complementary direction and how their work could fit into that narrative, were it to transpire to be successful at scale.", "essential_references_not_discussed": "N/A.", "other_strengths_and_weaknesses": "N/A.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "Please see the questions in my responses above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yfDIlXntmG", "reviewer": "Reviewer_dsES", "summary": "The work aims to develop a new model-based reinforcement learning method that trained on a set of tasks with consistent changes and different reward functions. The researchers assume that the environment's dynamics remain the same for all tasks. They use TD-MPC as the basic approach, and train a separate generative model (VAE) that simulates the generalized transition function using a dataset of the new task and synthetic data from the old version of the model. The current transition model is then trained using an error minimization loss function on the current dataset and error regularization on the old model's on data generated by the generative model. Additionally, the authors propose using internal rewards to encourage visits to states that are accurately predicted by both the old and current models.\nTo evaluate the effectiveness of their approach, the authors compare it with traditional TD-MPC and an older EWC baseline on two tasks: miniGrid and MuJoCo.\n\n## update after rebuttal\nBoth before and after the rebuttal phase, I believe that the work has a certain novelty, and therefore, I leave my current high assessment.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your very positive evaluation and accept recommendation. We believe we can resolve your concerns.\n\nQ: TDMPC2’s Multitask Training\n\nA:We clarify that TDMPC2 is evaluated in a multitask regime: it is trained on all tasks jointly, with access to the full replay buffers from every task simultaneously and (implicitly) the task identity or reward function for each experience. This means TDMPC2 benefits from seeing all task data at once (no task ordering) and can leverage task-specific information during training. In contrast, **DRAGO** tackles tasks in a strict continual learning manner: tasks are presented sequentially, and our method does not utilize task IDs or any access to past task data/replay buffers once those tasks are finished. DRAGO must retain knowledge through its mechanisms (generative rehearsal and intrinsic reward) without being able to directly revisit past data.\n\nQ: Additional Continual Learning Baselines\n\nA:We agree that incorporating more recent state-of-the-art baselines will strengthen the evaluation. In the original submission, we compared against naïve continual fine-tuning (Continual TDMPC), a from-scratch retraining baseline, and EWC. To address the reviewer’s suggestion, we are **running two new baselines** and will include them on this link https://drive.google.com/file/d/18TdI9nPCT7MMQCQMmL1ZYxJkblhUI91i/view?usp=sharing once we get the results: **(1) Pseudo-rehearsal with world models**: a generative replay baseline where an agent uses a fixed pretrained world model (or VAE) to rehearse past experiences, akin to DRAGO’s Synthetic Experience Rehearsal but without our intrinsic exploration component. **(2) Bounded replay-buffer baseline**: an experience replay approach that stores a limited number of real past transitions (capped to a similar memory size as our generative model) and interleaves them during training on new tasks. This will provide a direct yardstick for DRAGO’s no-data approach: i.e., how well does compressing experiences into a model compare to simply saving raw data with equal storage budget.\n\nQ: Curriculum Learning\n\nA:Thank you for the interesting suggestion on curriculum learning. We agree that intelligently ordering tasks (e.g., from easier to harder or with gradually increasing complexity) could further improve continual learning performance. A curriculum might help the agent build up its world model in a more structured way, potentially reducing forgetting and improving transfer. However, designing an optimal curriculum is non-trivial and was beyond the scope of our current work, which focuses on general mechanisms applicable to any task sequence. We opted to evaluate DRAGO on diverse task sequences without assuming a favorable order. Nonetheless, exploring curriculum learning in conjunction with DRAGO is an exciting avenue for future work. We will note this in the discussion as a potential enhancement, as it could complement our approach by easing the learning progression through tasks.\n\nQ: Handling Image Observations\n\nA: Thank you for the suggestion. We would like to first point out that the domains that we tested on (MIniGrid & Deepmind Control Suite) are two of the most popular RL benchmarks that have been evaluated in a large number of prior papers and have been shown to be quite challenging environments. The tasks we designed for evaluating transfer performance are even more challenging on DMC tasks as they require the agent to learn to transition from one locomotion mode (jump, run etc.) to another (run forward, run backward). While due to time constraints we cannot test the method on image observation settings, we think this is an exciting and challenging problem for future work as we mentioned at the end of Section 3.1, especially as we can replace VAE with diffusion models that are capable of generating high-quality image data.\n\nQ: Clarification of Notation $q_{\\theta_i}$ in Eq. (6)\n\nA: We apologize for the confusion regarding the notation $q_{\\theta_i}$ in Equation (6). In the context of our VAE-based generative model, $q_{\\theta_i}$ denotes the encoder’s approximate posterior distribution for task $i$. In other words, $q_{\\theta_i}(z \\mid s,a)$ is the VAE encoder’s output: the probability distribution (in latent space) that approximates the true posterior of the latent variable $z$ given an observation-action pair $(s,a)$. The subscript $i$ indicates that this encoder (with parameters $\\theta_i$) is the one learned up to task $i$ (since we train a new generative model $G_i$ for each task $i$ using both new and past data via rehearsal). We will explicitly clarify this in the revised paper. Essentially, Eq. (6) is the standard VAE loss: $L_{\\text{gen}}(\\theta_i) = E_{(s,a) \\sim D_{\\text{gen}}}\\Big[ -E_{z \\sim q_{\\theta_i}(z|s,a)}[\\log p_{\\theta_i}(s,a \\mid z)] + \\text{KL}(q_{\\theta_i}(z|s,a)|p(z)) \\Big]$, where $q_{\\theta_i}$ is the encoder’s distribution and $p_{\\theta_i}$ is the decoder (generative distribution).", "claims_and_evidence": "The main claims of the authors regarding the effectiveness of using the generative model for previous tasks are confirmed by experiments.", "methods_and_evaluation_criteria": "The method itself is relatively new, and while it mainly utilizes the TD-MPC code base, it does provide an improvement in the continuous learning process. I should note that TD-MPC also claims to have some form of multitasking ability and the capacity to work with multiple tasks at once. However, the authors do not elaborate on this in any detail. Instead, the authors use relatively simple vector control systems to create a series of tasks that need to be solved. The method in question is limited to vector envs only and is unlikely to be effective for observations presented in the form of images.", "theoretical_claims": "The paper provides a conclusion of the main loss function in equation 5, which can be considered sufficient theoretical justification for the proposed method.", "experimental_designs_or_analyses": "It should be noted that the test and training tasks are not randomized, and the authors perform experiments on fixed sequences of tasks. This raises questions about the generalizability of their results to other sequences and compositions of test tasks. Additionally, the authors use an old EWC continuous learning baseline, which is a significant drawback and does not reflect the current state of the field.", "relation_to_broader_scientific_literature": "The authors did not adapt or demonstrate the effectiveness of other continuous learning methods mentioned in the review in any way. Yes, the environment model is not used there, but perhaps this is not so necessary for such simple environments as Minigrid and MuJoCo.", "essential_references_not_discussed": "The authors have provided all the necessary references and methods.", "other_strengths_and_weaknesses": "It would also be useful to compare the authors' approach with different curriculum learning scenarios based on the same sequence of tasks. It would also be a good baseline.", "other_comments_or_suggestions": "The authors did not decipher the notation q_\\theta_i in formula 6.", "questions_for_authors": "To what extent would any curriculum learning scenarios also be effective for such task sequences?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Fa2TIMCK9v", "reviewer": "Reviewer_rfpF", "summary": "This paper presents a new approach to model-based continuous reinforcement learning (DRAGO) aimed at improving the incremental development of world models across a range of tasks. DRAGO consists of two key components: Synthetic Experience Rehearsal and Regaining Memories Through Exploration. Empirical evaluations have shown that DRAGO is capable of preserving knowledge across a wide range of tasks and achieving excellent performance in a variety of continuous learning scenarios.\n\n## update after rebuttal\nI have read the authors' rebuttal and appreciate their hard work and detailed responses. While their clarifications have helped me better understand certain points that were previously unclear, they are not sufficient for me to increase my rating score further.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive feedback and thoughtful questions.We address your questions below.\n\nQ: “how the synthetic experiences generated by the generative model differ from actual past experiences in terms of their representational accuracy?”\n\nA: Synthetic experiences in DRAGO are generated by a continually learned VAE-based generative model $G$ that **encodes and decodes both states and actions**, capturing the joint distribution of prior state-action. In other words, the agent “dreams” trajectories from its learned world model. These synthetic experiences **approximate real past interactions**—if $G$ is well-trained, the sampled states and actions resemble those from earlier tasks, helping reinforce previously learned dynamics **without storing actual data**. We acknowledge that synthetic data may not be perfect (due to model approximation error, we also discussed this in appendix F), but our design mitigates this: we **retrain $G$ after each task on a mix of new and generated old data** so it retains the ability to produce samples representative of all past tasks. Moreover, **DRAGO’s second component (Regaining Memories via Exploration)** complements generative rehearsal by actively revisiting important states in the real environment. This intrinsic reward-driven exploration addresses any gaps in the generative model’s coverage, ensuring that the world model doesn’t diverge from true environment dynamics.\n\n\nQ: Additional Continual Reinforcement Learning Baseline\n\nA: We agree that incorporating more recent state-of-the-art baselines will strengthen the evaluation. In the original submission, we compared against **naïve continual fine-tuning (Continual TDMPC)**, a **from-scratch retraining baseline**, and **EWC**. To address the reviewer’s suggestion, we are **running two new baselines** and will include them on this link https://drive.google.com/file/d/18TdI9nPCT7MMQCQMmL1ZYxJkblhUI91i/view?usp=sharing once we get the results: **(1) Pseudo-rehearsal with world models**: a generative replay baseline where an agent uses a learned world model (or VAE) to rehearse past tasks’ experiences, akin to DRAGO’s Synthetic Experience Rehearsal but without our intrinsic exploration component and continual learning of the generative model. **(2) Bounded replay-buffer baseline: an experience replay approach that stores a limited number of real past transitions (capped to a similar memory size as our generative model)** and interleaves them during training on new tasks. This will provide a direct yardstick for DRAGO’s no-data approach: i.e., how well does compressing experiences into a model compare to simply saving raw data with equal storage budget.\n\nQ: “How is the action sampled in continuous action space for synthetic data”\n\nA:  We sample actions jointly with states from the generative model, rather than picking random actions. This is crucial for continuous action domains: **random actions might not lead to meaningful or realistic transitions**, whereas our generative model produces plausible $(s, a)$ pairs grounded in past experience​.\n\nQ: **Details on Sampling from $p_G(s,a;\\theta)$ in Eq. (1)**\n\nA: In Equation (1) of the paper, $(\\hat{s}, \\hat{a}) \\sim p_G(s,a;\\theta)$ denotes drawing a state-action pair from the generative model’s distribution. As explained above, this is implemented by sampling from the VAE. We will clarify in the text that $p_G$ is the VAE-generative distribution over state-action pairs. The sampled pair $(\\hat{s}, \\hat{a})$ is then fed into the **frozen previous dynamics model $T_{\\text{old}}$** to predict the next state $\\hat{s}' = T_{\\text{old}}(\\hat{s}, \\hat{a})$. This yields a synthetic transition $(\\hat{s}, \\hat{a}, \\hat{s}')$ which is used (alongside real data from the new task) to train the current dynamics model. Thus, **the generative model provides realistic past states and actions, and $T_{\\text{old}}$ ensures the next-state is generated according to learned physics**.\n\nQ: Train from scratch approach’s performance\n\nA: You are correct that in a few scenarios (notably the Cheetah run-backward task), the scratch baseline slightly overtakes continual learning baseline. We observed this in our results: the continual learning method does **not fully eliminate plasticity loss**, and a sufficiently different new task can benefit from a fresh model. We hypothesize that in the backward-running tasks, the agent’s prior knowledge (largely acquired from forward-running dynamics) was less applicable or even somewhat **biasing the exploration**. We thank the reviewer for pointing this out, and we will clarify in the revision.", "claims_and_evidence": "YES", "methods_and_evaluation_criteria": "YES", "theoretical_claims": "I examined the theoretical proofs of two key components of the author's approach: Synthetic Experience Rehearsal and Regaining Memories Through Exploration. The logic is basically correct, but there is still a lack of clarity that needs to be improved.", "experimental_designs_or_analyses": "The experimental design is generally sound, but still insufficient, and additional supplementary experiments need to be added.", "relation_to_broader_scientific_literature": "The task addressed in this paper is a hot topic of discussion in recent years, where models learn new tasks with catastrophic forgetting of past tasks. This has been a great challenge in the field of continuous learning and reinforcement learning, and there have been many previous studies on continuous reinforcement learning (lifelong reinforcement learning). The authors' main innovation comes from a research finding on dreaming in the 1990s, which has been applied to other modelling domains, but not to the field of continuous reinforcement learning (lifelong reinforcement learning).", "essential_references_not_discussed": "No significant relevant papers were found that were not cited.", "other_strengths_and_weaknesses": "Strengths:\n1. The paper explains the innovativeness clearly enough, and the structure of the article and the linguistic description make it very easy for the reader to understand the proposed innovativeness.\n2. The presentation of previous relevant work is clearer and may be a continuation of the team's previous work.\n3. The algorithms and pseudocode are more adequately described and innovatively explained in the supplementary material.\n4. The introduction of the problem through the robot was very easy to understand\n\nWeaknesses:\n1. In Section 3.1, the concept of “Synthetic Experience Rehearsal” is introduced as a method for generating synthetic experiences from past tasks. Could the authors clarify how the synthetic experiences generated by the generative model differ from actual past experiences in terms of their representational accuracy?\n2. Too few algorithms for comparative testing. It is recommended that the authors add more SOTA methods on continuous reinforcement learning (lifelong reinforcement learning) published in top journals/conferences in recent years\n3. In Eq. (1), where the synthetic state is generated using the frozen old world model Told, could the authors provide more explanation about how the state-action pair is sampled from the generative model p_G(s, a; \\theta)? Specifically, how is the action sampled in continuous action spaces, and how do the learned transition dynamics handle this?\n4. It’s suggested to quantitatively demonstrate the degree of forgetting in the main text, e.g., by reporting how well each method recovers performance (or world model error) on earlier tasks after the training sequence has ended. While Table 3 in the Appendix of the paper provides a comparison of the final performance of DRAGO and the baseline on each of the training tasks, it would have been more convincing in the main paper to mention that ‘DRAGO roughly maintains the performance of the old tasks without degrading the performance of the new tasks’.\n5. On some of the test tasks, the ‘train from scratch’ approach came close to or even outperformed the continuous learning baseline (for the ‘Cheetah backward’ related combinatorial task). It is suggested that the authors briefly analyse the reasons for this phenomenon.\n6. There still seems to be a clerical error in the title and abbreviation.", "other_comments_or_suggestions": "N/A.", "questions_for_authors": "N/A.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xzKoNfeWmQ", "reviewer": "Reviewer_XZ8k", "summary": "The authors propose a method for continual model-based reinforcement learning, where, ideally, an agent retains previously learned skills while learning new skills mitigating the catastrophic forgetting problem. The main problem addressed is the bounded storage problem, where the agent is not assumed to have infinite storage for storing previous transitions, which most continual RL methods assume (replay-based approaches). Instead, this work consists of two key components:\n1. Synthetic Experience Rehearsal, where a generative model (VAE) produces synthetic transitions from all previous tasks, which are used in tandem with the current tasks data to train the new transition model.\n2. Regaining Memories Through Exploration, which is an intrinsic reward that encourages exploration of states that were well-understood in previous tasks but are not currently well understood. This supposedly helps to bridge newly learned tasks and previously learned tasks.\n\nHence, instead of storing previous transitions, the proposed method stores a generative model that learn the distribution of the previous transitions. Thus, the generative model should require less storage than keeping previous transitions.\nThe experimental evidence is conducted through continual learning experiments on the MiniGrid environment and the Deepmind Control Suite. The method is compared to the following baselines: TDMPC from scratch (model-based/world-model RL), 'continual' TDMPC (initialize with model of previous task), EWC (regularization-based continual RL).\n\nThe presented results significantly outperform the baselines in most experiments.\n\nIn an ablation study, the authors also verify the complementary roles of the two following components:\n1. The generative model\n2. The intrinsic reward\n\nWhile I provided a \"Reject\" as score, I think that this paper overall details valuable research, that will be accepted if the authors improve the clarity and above all justify the use of the Generative Model to store experience through experimental support (comparing with CRL baselines).", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed and constructive feedback. We are encouraged that you find our ideas valuable and appreciate your suggestions on clarity and comparisons. We address each concern below:\n\nQ: Replay-based Baseline\n\nA: Thank you for pointing this out – we are adding a **fixed-size replay buffer** baseline with the same storage budget as DRAGO’s generative model. Preliminary implementation is underway, and we will share results on this anonymous link https://drive.google.com/file/d/18TdI9nPCT7MMQCQMmL1ZYxJkblhUI91i/view?usp=sharing after we have the results. This baseline is an experience replay approach that stores a limited number of real past transitions (capped to a similar memory size as our generative model) and interleaves them during training on new tasks. Our hypothesis is that DRAGO will perform on par with or better than this bounded replay baseline, especially as tasks accumulate (since DRAGO can simulate diverse past experiences without strictly limited slots).\n\nQ: Pseudo-rehearsal Baseline\n\nA: Since the original paper does not provide source code, we are currently in the process of implementing a **world-model pseudo-rehearsal** baseline based on the paper. This is a generative replay baseline where an agent uses a pretrained world model (or VAE) to rehearse past tasks’ experiences, akin to DRAGO’s Synthetic Experience Rehearsal but without our intrinsic exploration component and continual learning of the generative model. We will share results on this anonymous link https://drive.google.com/file/d/18TdI9nPCT7MMQCQMmL1ZYxJkblhUI91i/view once the experiments are finished.\n\nQ: Terminology “Continual TDMPC”: \n\nA: Thank you for flagging this potential misunderstanding. In our paper, “Continual TDMPC” refers to the naïve baseline where we initialize the model for each new task with the previously learned TDMPC world model, without any forgetting mitigation. It is essentially the standard TDMPC algorithm simply fine-tuned sequentially.\n\nQ: Paper structure and readability:\n\nA: We appreciate these suggestions and will incorporate them in the final version. We will add a bulleted summary of contributions in the introduction. Specifically,\n\n- **Novel Continual MBRL Framework**. We introduce **DRAGO**, a new approach for continual model-based reinforcement learning that addresses catastrophic forgetting while incrementally learning a world model across sequential tasks without retaining any past data.\n\n- **Synthetic Experience Rehearsal**. We propose a generative replay mechanism that synthesizes “old” transitions using a learned generative model alongside a frozen copy of the previously trained world model. This synthetic data consistently reinforces earlier dynamics knowledge, mitigating forgetting in each new task.\n\n- **Regaining Memories Through Exploration**. We design an intrinsic reward signal that nudges the agent toward revisiting states that the old model explained well—effectively “reconnecting” current experiences with previously learned transitions. This mechanism complements the synthetic rehearsal by incorporating real environmental interactions to maintain a more complete world model.\n\n- **Extensive Empirical Validation**. Through experiments on **MiniGrid** and **DeepMind Control Suite** domains, we show that DRAGO:\n\n1. Substantially improves knowledge retention compared to standard continual MBRL baselines.\n\n2. Achieves higher forward-transfer performance, allowing faster adaptation to entirely new (but related) tasks.\n\n3. Exhibits strong few-shot learning capabilities, substantially outperforming both learning-from-scratch and other continual methods under limited interaction budgets.\n\nQ: Formalism clarity (transition model, loss equations, notation)\n\nA: We apologize for the confusion. In our notation, $T$ is the parametric transition model (the learned dynamics predictor) and $p(\\cdot)$ denotes a probability or distribution. For example, $p(s' \\mid s,a;\\psi)$ is the likelihood of observing $s'$ given $(s,a)$ under the transition model $T_{\\psi}$. We will explicitly define the transition model in the main text upon first use and use a consistent notation (e.g. using $T$ for the model and $p$ for probabilities) throughout. We will also clarify the loss equations step-by-step. In particular, Equation (4) in the paper decomposes the loss into current-task loss (on real data $D_i$) and synthetic rehearsal loss (on generated data $\\hat D$). Equation (5) then shows the combined training objective for the dynamics model: it includes a term $|T_i(s,a)-s'|^2$ for real transitions and a term $|T_{i}(ŝ,â)-T_{old}(ŝ,â)|^2$ for synthetic ones (weighted by $\\lambda$). We will make sure to clearly explain each term and the roles of $T$ vs. $p$ in a revision.\n\nQ: Minor issues (typos, notation, algorithm pseudocode)\n\nA:  We will fix all minor typos and notational errors. We also agree that a high-level pseudocode algorithm in the main paper would aid clarity.", "claims_and_evidence": "The claims are not clearly stated, but I extracted the following.\n- New tasks are learned without forgetting previous tasks.\n\t- Evidence: Supported by experiment on gridworld (nicely visualized in Figure 4)\n- The previously learned skills improve few-shot transfer to new tasks that build upon the previous skills.\n\t- Evidence: supported by experiments on cheetah/walker\n- Synthetic Experience Rehearsal learns the distribution of previous experiences\n\t- Evidence: partly shown in Figure4.\n- The intrinsic reward incentivices agent to relearn previously learned tasks and bridges the gap\n\t- Evidence: only in ablation study that shows worse performance without it.\n- Mitigating the bounded storage problem\n\t- Evidence: No evidence on environments with a large number of tasks, where replay-based continual RL would be limited.\n\t\t- Or alternatively: comparison to replay-based methods and how much less storage is needed for similar performance", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria make sense for the problem at hand.\n\nHowever, the evaluation would benefit from another experiment demonstrating good performance where replay-based methods fail to verify the main claim of solving the bounded storage problem. Specifically, as the authors are reusing the technique of learning a generative model instead of storing the raw replay buffer, they should demonstrate the extent to which this technique is more efficient than storing a replay buffer. Thus, the authors could use replay buffers with limited capacity (to match the storage requirements of their generative model). Then, uniformly drawing could constitute a first baseline, even if more advance storage selection and sampling have been developed in CL papers.", "theoretical_claims": "Not applicable", "experimental_designs_or_analyses": "The paper lacks experiments designed to test the performances on an environment with a larger number of tasks, which is especially relevant if the main goal is to be an alternative to replay-based methods in settings where they become infeasible due to a large number of previous transitions (bounded storage assumption). Could you extend the number of rooms to e.g. 16 ?\nOn a similar note, the method is not compared to replay-based methods, which would probably perform better or similar on the conducted experiments (according to my CL expert fellows).\n\nMaybe a valuable experimental question would thus be: \"how close can we match their performance with x less storage use?\"\n\nTD-MPC baselines: Please clarify how exactly 'from scratch' and continual TDMPC differ. \nDoes the continual TDMPC use the same task-embedding as the model from the previous task?\n\nSome baselines are missing. For example: pseudo-rehearsal (Atkinson et al.) and world-model pseudo-rehearsal (Ketz et al.) -> see below.", "relation_to_broader_scientific_literature": "This research work combines known methods from continual learning (pseudo-recursal) with methods from reinforcement learning (MBRL).\nSome essential references are not discussed:\n- Pseudo-recursal is not referenced, neither is pseudo-rehearsal.\n\t- which is exactly the idea of [1], just with DQN instead of MBRL\nThere even exists methods applying pseudo-rehearsal to MBRL [2] (but not published) \n\nHow does your method differ? You should explain why these baselines are excluded, or if they are not relevant to support your claims (but then the reader would need explicit claims listed)\n\n[1] Atkinson, Craig, et al. \"Pseudo-rehearsal: Achieving deep reinforcement learning without catastrophic forgetting.\" *Neurocomputing* 428 (2021): 291-307.\n[2] Ketz, N., Kolouri, S., & Pilly, P. (2019). Continual learning using world models for pseudo-rehearsal. arXiv preprint arXiv:1903.02647.", "essential_references_not_discussed": "See above", "other_strengths_and_weaknesses": "Strength:\n- Figure 4 is really nice, although a short claim as first sentence what this figure shows would be great.\nIn general you could improve the captions of all figure following this principle:\n   - The first sentence should highlight the main message of the Figure/Table. E.g. **DRAGO reduce the catastrophic forgetting from previous tasks,** \n   - The next sentences then explain what is depicted in the Table/Figure, E.g Illustrated by the prediction score of the learned world models across the entire gridworld after each task.... etc\n\nWeaknesses:\n- Missing clear description of all research questions answered / main contributions.\n- Figures:\n\t- architecture overview\n\t- do not always contain a direct description of what to take away (as explained above)\n\t- Figure 3 does not add a lot -> remove?\n\t- Figures 5,6 could use a grid\n\t- missing number of seeds/runs (also at table 1)\n- Formalism:\n\t- TDMPC2 uses the MSE on the embedded state. Are the states embedded in DRAGO?\n\t\t- If not: This is not really TDMPC -> clarify in method\n\t\t- If yes: not clear from context, redefine s.\n\t- The transition model $T_i$ and it's probability are never properly defined. What kind of model is it?\n\t- Makes it sound like its two different models (also with eq.4 and 5 using one each)\n\t- Suggestion: Just use $p_T(s'|s,a;\\psi)$ and $\\mathbb{E}[p_T(s'|s,a;\\psi)]$ , or define T that way.\n\t- Loss formulation is confusing, because you show the total loss using probabilities and then split it later into dyn and gen loss, whereas the dyn loss is actually MSE that uses the deterministic(?) T.\n\t- Also, the total loss (eq. 4) does not specify, where $\\hat{s}'$ comes from, which is only later revealed to be from $T_\\text{old}$ .\n\t\t- Maybe equations 3 and 4 are unnecessary (and confusing) and can be ommited altogether?\n\t\t- Instead just show L_dyn and L_gen and then: L_total = L_dyn + L_gen\n- Baselines:\n\t- Continual TDMPC is not really continual, more like 'pre-initialized'\n\t- Missing comparison to a replay-based method to show sota results and how drago compares.\n- Limitations missing in main paper. Super important, since the generated synthetic state-action pairs may not fully represent all previous tasks with a growing number of tasks, as discussed in Appendix F.\n- A (short) high-level algorithm in the main paper would also be beneficial.\n\nMore specifically, to improve clarity, I strongly advise the authors to include in their paper:\nI. A list of contributions at the end of the introduction (often denoted with i., ii., ...) that stands out visually.\nII. A list of scientific questions at the beginning of the experimental evaluation section (often denoted Q1, Q2, ... etc), that are each answered in different paragraphs. For example\nQ1/ How does DRAGO perform in comparison to existing CRL baselines?\nQ2/ How much does the generative model help preventing catastrophic forgetting?\nQ3/ Is the generative model more storage efficient than the classical use of a replay buffer?\n...", "other_comments_or_suggestions": "Background formalism MDP tuple: Usually use <> instead of ()\nTypos:\n- Chapter title 3\n- Lines 4-7 under eq. 5\n- Just before chapter 4: 2x same sentence.\n\nWhile I graded the paper with a reject. I think that the paper does provide a valuable research work, that, once improved, will be of valuable interest for many readers. However, in its current form, it cannot be accepted to a major conference such as ICML. I hope that the authors can use the provided feedback to improve their work.", "questions_for_authors": "- Q1: add research questions / key contributions\n- Q2: properly define transition model/likelihood, also focus on one notation (T or p)\n- Q3: reformulate loss definitions for clarity (see above)\n- Q4: add limitations + high-level algorithm\n- Q5: add replay-based baseline -> how does DRAGO improve upon that? (e.g. how much % less storage needed for similar performance))\n- Q6: add pseudo-rehearsal baselines (see references above)\n- Q7: benchmark number of tasks -> when does it fail?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "Not applicable"}}]

Author rebuttal:

Meta-review: This paper proposes an algorithm built on top of TDMCP which learns a generative model of prior environment experience and an exploration objective based on revisiting prior memories, for the ultimate objective on improving on catastrophic forgetting in continual learning settings. This is an important setting and the contribution here seems well-motivated and well-justified through experimentation: limited memory (i.e., as opposed to doing pure in-context learning) and privacy are important aspects to touch on, and the results are strong.

There were concerns on benchmarks, for example the environments being too limited or cherry picked and the baselines being too limited, but these seem to have been addressed in the rebuttal. There also seemed to be various concerns about clarity, but from reading the responses I believe these were sufficiently addressed, given that the final manuscript contains these clarifications.

I therefore recommend acceptance to the conference.
Decision: Accept (poster)