Distilling Reinforcement Learning into Single-Batch Datasets

TL;DR

This paper introduces a method to compress reinforcement learning environments into small, synthetic supervised datasets, enabling rapid training and cross-modality learning from RL to supervised learning.

Contribution

It extends dataset distillation to reinforcement learning, transforming RL tasks into one-batch supervised datasets and demonstrating its generalizability across environments and architectures.

Findings

RL environments can be distilled into single-batch datasets

Distilled datasets enable one-step training approximating original RL performance

Method generalizes across different RL environments and learner architectures

Abstract

Dataset distillation compresses a large dataset into a small synthetic dataset such that learning on the synthetic dataset approximates learning on the original. Training on the distilled dataset can be performed in as little as one step of gradient descent. We demonstrate that distillation is generalizable to different tasks by distilling reinforcement learning environments into one-batch supervised learning datasets. This demonstrates not only distillation's ability to compress a reinforcement learning task but also its ability to transform one learning modality (reinforcement learning) into another (supervised learning). We present a novel extension of proximal policy optimization for meta-learning and use it in distillation of a multi-dimensional extension of the classic cart-pole problem, all MuJoCo environments, and several Atari games. We demonstrate distillation's ability to compress complex RL environments into one-step supervised learning, explore RL distillation's generalizability across learner architectures, and demonstrate distilling an environment into the smallest-possible synthetic dataset.