Table of Contents
Fetching ...

Continual Policy Distillation from Distributed Reinforcement Learning Teachers

Yuxuan Li, Qijun He, Mingqi Yuan, Wen-Tse Chen, Jeff Schneider, Jiayu Chen

TL;DR

This work tackles continual reinforcement learning by decoupling it into parallel single-task RL training (teachers) and a central generalist model that distills from these teachers. It introduces a Transformer-based Mixture-of-Experts central model with incremental expert expansion, a contrastive task embedding to infer latent task context, and a hybrid forgetting mitigation combining diversity-aware trajectory replay with parameter masking. The approach demonstrates strong plasticity and stability on Meta-World MT25, recovering the majority of teacher performance (well over 85%) while keeping task-wise forgetting under 10%, and shows robustness under multi-stage distillation. These results suggest a scalable path toward generalist embodied agents by transforming continual RL challenges into supervised-like distillation dynamics, with implications for efficient pretraining and lifelong learning in robotics and related domains.

Abstract

Continual Reinforcement Learning (CRL) aims to develop lifelong learning agents to continuously acquire knowledge across diverse tasks while mitigating catastrophic forgetting. This requires efficiently managing the stability-plasticity dilemma and leveraging prior experience to rapidly generalize to novel tasks. While various enhancement strategies for both aspects have been proposed, achieving scalable performance by directly applying RL to sequential task streams remains challenging. In this paper, we propose a novel teacher-student framework that decouples CRL into two independent processes: training single-task teacher models through distributed RL and continually distilling them into a central generalist model. This design is motivated by the observation that RL excels at solving single tasks, while policy distillation -- a relatively stable supervised learning process -- is well aligned with large foundation models and multi-task learning. Moreover, a mixture-of-experts (MoE) architecture and a replay-based approach are employed to enhance the plasticity and stability of the continual policy distillation process. Extensive experiments on the Meta-World benchmark demonstrate that our framework enables efficient continual RL, recovering over 85% of teacher performance while constraining task-wise forgetting to within 10%.

Continual Policy Distillation from Distributed Reinforcement Learning Teachers

TL;DR

This work tackles continual reinforcement learning by decoupling it into parallel single-task RL training (teachers) and a central generalist model that distills from these teachers. It introduces a Transformer-based Mixture-of-Experts central model with incremental expert expansion, a contrastive task embedding to infer latent task context, and a hybrid forgetting mitigation combining diversity-aware trajectory replay with parameter masking. The approach demonstrates strong plasticity and stability on Meta-World MT25, recovering the majority of teacher performance (well over 85%) while keeping task-wise forgetting under 10%, and shows robustness under multi-stage distillation. These results suggest a scalable path toward generalist embodied agents by transforming continual RL challenges into supervised-like distillation dynamics, with implications for efficient pretraining and lifelong learning in robotics and related domains.

Abstract

Continual Reinforcement Learning (CRL) aims to develop lifelong learning agents to continuously acquire knowledge across diverse tasks while mitigating catastrophic forgetting. This requires efficiently managing the stability-plasticity dilemma and leveraging prior experience to rapidly generalize to novel tasks. While various enhancement strategies for both aspects have been proposed, achieving scalable performance by directly applying RL to sequential task streams remains challenging. In this paper, we propose a novel teacher-student framework that decouples CRL into two independent processes: training single-task teacher models through distributed RL and continually distilling them into a central generalist model. This design is motivated by the observation that RL excels at solving single tasks, while policy distillation -- a relatively stable supervised learning process -- is well aligned with large foundation models and multi-task learning. Moreover, a mixture-of-experts (MoE) architecture and a replay-based approach are employed to enhance the plasticity and stability of the continual policy distillation process. Extensive experiments on the Meta-World benchmark demonstrate that our framework enables efficient continual RL, recovering over 85% of teacher performance while constraining task-wise forgetting to within 10%.
Paper Structure (35 sections, 7 equations, 8 figures, 9 tables, 1 algorithm)

This paper contains 35 sections, 7 equations, 8 figures, 9 tables, 1 algorithm.

Figures (8)

  • Figure 1: We propose a teacher-student framework where single-task experts are initially trained on the Meta-World benchmark to generate demonstration trajectories. These trajectories are then used to distill a central, Transformer-based student model with a Mixture-of-Experts (MoE) architecture. We freeze the shared parameters (e.g., embeddings, attention) after the first stage, while expanding the model's capacity by adding one new expert to every layer per stage.
  • Figure 2: Two main steps explaining how continual policy distillation helps maintain the stability of the central model: 1) expert expansion and parameter-wise masking to help retain memory; 2) replay-based approach to mitigate forgetting. Less than 10% of the distill data is selected through the DPP algorithm as the replay buffer.
  • Figure 3: PCA projection of task embeddings on Meta-World MT25. The visualization highlights representative pairs of tasks that share similar manipulation primitives and reveals that the learned embeddings cluster according to manipulation mechanics (e.g., Button Press group, Pick Place group), validating the effectiveness of the contrastive objective in capturing task similarity.
  • Figure 4: Joint distillation results: (a) using different algorithms for replay trajectory selection; (b) selecting a different number of trajectories for the replay buffer.
  • Figure 5: Comparative analysis of continual policy distillation performance. (a) Overall performance; (b) Backward transfer (BWT). Our method outperforms all other baselines on average accuracy and backward transfer.
  • ...and 3 more figures