Policy Diffusion: Generative Models for Policy Datasets

Recent progress in Quality Diversity Reinforcement Learning (QD-RL) has enabled learning a collection of behaviorally diverse, high performing policies. However, these methods typically involve storing thousands of policies, which results in high space-complexity and poor scaling to additional behaviors. An alluring proposition would be to condense the archive into a single model while retaining the performance and coverage of the original collection of policies, which has proved challenging. In this work, we propose using diffusion models to distill the archive into a single generative model over policy parameters. We show that our method achievesa compression ratio of 13x while recovering 98% of the original rewards and 89% of the original coverage. Furthermore, the conditioning mechanism of diffusion models allows for flexibly selecting and sequencing behaviors using language.