Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Augmenting Replay in World Models for Continual Reinforcement Learning

Luke Yang, Levin Kuhlmann, Gideon Kowadlo

TL;DR

The paper tackles continual reinforcement learning by combining model-based world models with a memory-efficient augmented replay buffer. It extends DreamerV3 to WMAR, introducing a dual-buffer LTDM mechanism and spliced rollouts to mitigate forgetting while preserving plasticity, evaluated on Procgen and Atari benchmarks. Results show substantial improvements in forgetting and competitive forward transfer, especially in tasks without shared structure, with modest gains in shared-structure settings. This work demonstrates that model-based RL with compact, distribution-matching replay can effectively support continual learning and motivates further exploration of memory-efficient world-model replay strategies.

Abstract

Continual RL requires an agent to learn new tasks without forgetting previous ones, while improving on both past and future tasks. The most common approaches use model-free algorithms and replay buffers can help to mitigate catastrophic forgetting, but often struggle with scalability due to large memory requirements. Biologically inspired replay suggests replay to a world model, aligning with model-based RL; as opposed to the common setting of replay in model-free algorithms. Model-based RL offers benefits for continual RL by leveraging knowledge of the environment, independent of policy. We introduce WMAR (World Models with Augmented Replay), a model-based RL algorithm with a memory-efficient distribution-matching replay buffer. WMAR extends the well known DreamerV3 algorithm, which employs a simple FIFO buffer and was not tested in continual RL. We evaluated WMAR and DreamerV3, with the same-size replay buffers. They were tested on two scenarios: tasks with shared structure using OpenAI Procgen and tasks without shared structure using the Atari benchmark. WMAR demonstrated favourable properties for continual RL considering metrics for forgetting as well as skill transfer on past and future tasks. Compared to DreamerV3, WMAR showed slight benefits in tasks with shared structure and substantially better forgetting characteristics on tasks without shared structure. Our results suggest that model-based RL with a memory-efficient replay buffer can be an effective approach to continual RL, justifying further research.

Augmenting Replay in World Models for Continual Reinforcement Learning

TL;DR

The paper tackles continual reinforcement learning by combining model-based world models with a memory-efficient augmented replay buffer. It extends DreamerV3 to WMAR, introducing a dual-buffer LTDM mechanism and spliced rollouts to mitigate forgetting while preserving plasticity, evaluated on Procgen and Atari benchmarks. Results show substantial improvements in forgetting and competitive forward transfer, especially in tasks without shared structure, with modest gains in shared-structure settings. This work demonstrates that model-based RL with compact, distribution-matching replay can effectively support continual learning and motivates further exploration of memory-efficient world-model replay strategies.

Abstract

Continual RL requires an agent to learn new tasks without forgetting previous ones, while improving on both past and future tasks. The most common approaches use model-free algorithms and replay buffers can help to mitigate catastrophic forgetting, but often struggle with scalability due to large memory requirements. Biologically inspired replay suggests replay to a world model, aligning with model-based RL; as opposed to the common setting of replay in model-free algorithms. Model-based RL offers benefits for continual RL by leveraging knowledge of the environment, independent of policy. We introduce WMAR (World Models with Augmented Replay), a model-based RL algorithm with a memory-efficient distribution-matching replay buffer. WMAR extends the well known DreamerV3 algorithm, which employs a simple FIFO buffer and was not tested in continual RL. We evaluated WMAR and DreamerV3, with the same-size replay buffers. They were tested on two scenarios: tasks with shared structure using OpenAI Procgen and tasks without shared structure using the Atari benchmark. WMAR demonstrated favourable properties for continual RL considering metrics for forgetting as well as skill transfer on past and future tasks. Compared to DreamerV3, WMAR showed slight benefits in tasks with shared structure and substantially better forgetting characteristics on tasks without shared structure. Our results suggest that model-based RL with a memory-efficient replay buffer can be an effective approach to continual RL, justifying further research.
Paper Structure (30 sections, 3 equations, 4 figures, 2 tables)

This paper contains 30 sections, 3 equations, 4 figures, 2 tables.

Table of Contents

  1. Introduction
  2. World models with augmented replay (WMAR)
  3. World model
  4. Actor critic
  5. Augmented replay buffer
  6. Short-term FIFO buffer
  7. Long-term global distribution matching (LTDM) buffer
  8. Spliced rollouts
  9. Task-agnostic exploration
  10. Experiments
  11. Dataset
  12. Baselines
  13. Evaluation metrics
  14. Normalised rewards
  15. Forgetting (Backward transfer)
  16. ...and 15 more sections

Figures (4)

  • Figure 1: World Model overview and RSSM details.
  • Figure 2: The use of replay within the world model context. LTDM refers to the long-term distribution matching buffer.
  • Figure 3: Tasks without shared structure. The line is the median and the shaded area is between the 0.25 and 0.75 quartiles, of 5 seeds. Bold line segments show when a task is being trained. Scores are normalised by eq. \ref{['eq:normalization']}.
  • Figure 4: Tasks with shared structure. The line is the median and the shaded area is between the 0.25 and 0.75 quartiles, of 5 seeds. Bold line segments show when a task is being trained. Scores are normalised by eq. \ref{['eq:normalization']}.