Augmenting Replay in World Models for Continual Reinforcement Learning

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.

• The paper introduces WMAR which integrates a world model with an augmented replay buffer, dramatically reducing memory overhead and mitigating catastrophic forgetting.
• The method effectively improves forward transfer in tasks with shared structure while balancing performance retention in dissimilar tasks.
• Experimental results on OpenAI ProcGen and Atari benchmarks demonstrate WMAR’s stability and transfer efficiency compared to FIFO-only replay models.

Augmenting Replay in World Models for Continual Reinforcement Learning

Introduction to Continual Reinforcement Learning (CRL)

Continual reinforcement learning (CRL) is a challenging aspect of machine learning where an agent encounters a sequence of tasks, aiming to learn new tasks without losing acquired knowledge and enhance its performance on both previous and future tasks. Traditional reinforcement learning (RL) usually involves model-free approaches heavily reliant on large replay buffers to mitigate catastrophic forgetting. However, these approaches become inefficient due to the vast storage requirements. An alternative involves leveraging model-based RL using world models, aligning with the biological inspiration of replay mechanisms observed in mammalian brains. This work introduces World Models with Augmented Replay (WMAR) to address continual RL through an efficient memory system.

World Models with Augmented Replay (WMAR) Framework

WMAR integrates DreamerV3's world model framework with an augmented replay buffer, facilitating efficient memory usage in continual RL. The inherent architecture comprises a recurrent state-space model that simulates the environment's dynamics, reducing the reliance on excessive replay storage. A significant contribution of WMAR lies in the introduction of a memory-efficient distribution matching replay buffer alongside a short-term FIFO buffer for comprehensive experience training.

WMAR addresses two primary task settings:

• Tasks with Shared Structure: Leveraging a commonality between tasks allows sharing of learned knowledge across similar tasks.
• Tasks Without Shared Structure: Highlights the challenge of retaining performance in dissimilar task sets.

Figure 1: Continual learning performance on tasks without shared structure. Bold line segments denote periods in which certain tasks are trained. Scores are normalized by Equation~\ref{eq:normalization}.

Experimental Evaluation

The experimental framework evaluates WMAR against its FIFO-only variant across tasks with and without shared structure using benchmarks from OpenAI ProcGen and Atari environments. The primary metrics for evaluation include stability, backward/forward transfer, and forgetting.

The key observations from the experiments are as follows:

• Memory Efficiency: WMAR significantly reduces computational overhead compared to its FIFO-only counterpart while demonstrating improved performance retention in both shared and non-shared task settings.
• Forgetting and Transfer: WMAR showcases improved mitigation of catastrophic forgetting, especially in dissimilar task sets. However, this comes at the cost of reduced plasticity, impacting learning efficiency in new tasks.
• Performance Comparison: On tasks with shared structure, WMAR exhibits positive forward transfer, substantially enhancing learning speed and performance consistency across tasks.

  Figure 2: Continual learning performance on tasks with shared structure. Bold line segments denote the periods in which certain tasks are trained. Scores are normalized by Equation~\ref{eq:normalization}.

Implications and Future Outlook

WMAR demonstrates the potential of leveraging model-based RL with augmented replay buffers to address the continual learning problem efficiently. The reduction in memory usage without compromising performance emphasizes the viability of world model approaches for real-world continual RL applications. However, there is a need for further exploration into hyperparameter tuning and combining WMAR with other established techniques like behaviour cloning to address plasticity drawbacks.

In future research, it will be crucial to explore the intricacies of tuning world models to accommodate task variations effectively, particularly addressing the trade-off between stability and plasticity. Moreover, exploring the intersection of supervised tasks with continual learning paradigms may provide significant insights into the broader applicability of WMAR in diverse domains.

Conclusion

The introduction of WMAR marks a notable advancement in continual RL, illustrating the benefits of memory-efficient world models. The exploration of tasks with and without shared structures elucidates the complex dynamics of continual learning in model-based RL. The research evidences substantial progress but highlights avenues for further investigation into efficient model-based solutions for evolving task environments in reinforcement learning. The potential integration of WMAR with complementary RL strategies promises an exciting direction for enhancing agent adaptability and learning robustness.