Stable Continual Reinforcement Learning via Diffusion-based Trajectory Replay

TL;DR

This paper introduces DISTR, a diffusion-based trajectory replay method for continual reinforcement learning, which effectively mitigates catastrophic forgetting by memorizing high-return trajectories and prioritizing pivotal tasks, outperforming existing methods.

Contribution

The paper proposes a novel diffusion model-based replay mechanism for continual RL, addressing generative replay limitations and enhancing stability and plasticity in learning multiple tasks.

Findings

DISTR outperforms existing continual RL baselines on the Continual World benchmark.

The diffusion-based replay effectively preserves task knowledge and improves success rates.

Prioritization of pivotal tasks enhances learning efficiency and stability.

Abstract

Given the inherent non-stationarity prevalent in real-world applications, continual Reinforcement Learning (RL) aims to equip the agent with the capability to address a series of sequentially presented decision-making tasks. Within this problem setting, a pivotal challenge revolves around \textit{catastrophic forgetting} issue, wherein the agent is prone to effortlessly erode the decisional knowledge associated with past encountered tasks when learning the new one. In recent progresses, the \textit{generative replay} methods have showcased substantial potential by employing generative models to replay data distribution of past tasks. Compared to storing the data from past tasks directly, this category of methods circumvents the growing storage overhead and possible data privacy concerns. However, constrained by the expressive capacity of generative models, existing \textit{generative replay} methods face challenges in faithfully reconstructing the data distribution of past tasks, particularly in scenarios with a myriad of tasks or high-dimensional data. Inspired by the success of diffusion models in various generative tasks, this paper introduces a novel continual RL algorithm DISTR (Diffusion-based Trajectory Replay) that employs a diffusion model to memorize the high-return trajectory distribution of each encountered task and wakeups these distributions during the policy learning on new tasks. Besides, considering the impracticality of replaying all past data each time, a prioritization mechanism is proposed to prioritize the trajectory replay of pivotal tasks in our method. Empirical experiments on the popular continual RL benchmark \texttt{Continual World} demonstrate that our proposed method obtains a favorable balance between \textit{stability} and \textit{plasticity}, surpassing various existing continual RL baselines in average success rate.