Knowledge Retention in Continual Model-Based Reinforcement Learning

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:

Q: Replay-based Baseline

A: 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).

Q: Pseudo-rehearsal Baseline

A: 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.

Q: Terminology “Continual TDMPC”:

A: 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.

Q: Paper structure and readability:

A: We appreciate these suggestions and will incorporate them in the final version. We will add a bulleted summary of contributions in the introduction. Specifically,

• 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.

• 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.

• 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.

• Extensive Empirical Validation. Through experiments on MiniGrid and DeepMind Control Suite domains, we show that DRAGO:

1. Substantially improves knowledge retention compared to standard continual MBRL baselines.

2. Achieves higher forward-transfer performance, allowing faster adaptation to entirely new (but related) tasks.

3. Exhibits strong few-shot learning capabilities, substantially outperforming both learning-from-scratch and other continual methods under limited interaction budgets.

Q: Formalism clarity (transition model, loss equations, notation)

A: 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}(ŝ,â)-T_{old}(ŝ,â)|^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.

Q: Minor issues (typos, notation, algorithm pseudocode)

A: 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.

We thank the reviewer for their positive feedback and thoughtful questions.We address your questions below.

Q: “how the synthetic experiences generated by the generative model differ from actual past experiences in terms of their representational accuracy?”

A: 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.

Q: Additional Continual Reinforcement Learning Baseline

A: 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.

Q: “How is the action sampled in continuous action space for synthetic data”

A: 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​.

Q: Details on Sampling from $p_G(s,a;\theta)$ in Eq. (1)

A: 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.

Q: Train from scratch approach’s performance

A: 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.

Thank you for your very positive evaluation and accept recommendation. We believe we can resolve your concerns.

Q: TDMPC2’s Multitask Training

A: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.

Q: Additional Continual Learning Baselines

A: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.

Q: Curriculum Learning

A: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.

Q: Handling Image Observations

A: 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.

Q: Clarification of Notation $q_{\theta_i}$ in Eq. (6)

A: 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).

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:

Q: bounded memory and On-device MBRL

A: 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.

Q: Evaluation Procedure

A: 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.

Q: Integration with TD-MPC, Decoder, and Learner vs. Reviewer Agents

A: 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).

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.

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.