State Diversity Matters in Offline Behavior Distillation

TL;DR

This paper demonstrates that emphasizing state diversity in Offline Behavior Distillation improves policy performance, especially with limited original dataset diversity, by analyzing the roles of state quality and diversity.

Contribution

It uncovers the importance of state diversity in OBD, provides theoretical insights into error reduction, and introduces a simple density-weighted algorithm to enhance distillation.

Findings

State diversity correlates with better policy performance under high training loss.

The proposed SDW algorithm improves distillation results on datasets with limited diversity.

Theoretical analysis highlights the dominance of surrounding error over pivotal error in certain scenarios.

Abstract

Offline Behavior Distillation (OBD), which condenses massive offline RL data into a compact synthetic behavioral dataset, offers a promising approach for efficient policy training and can be applied across various downstream RL tasks. In this paper, we uncover a misalignment between original and distilled datasets, observing that a high-quality original dataset does not necessarily yield a superior synthetic dataset. Through an empirical analysis of policy performance under varying levels of training loss, we show that datasets with greater state diversity outperforms those with higher state quality when training loss is substantial, as is often the case in OBD, whereas the relationship reverses under minimal loss, which contributes to the misalignment. By associating state quality and diversity in reducing pivotal and surrounding error, respectively, our theoretical analysis establishes that surrounding error plays a more crucial role in policy performance when pivotal error is large, thereby highlighting the importance of state diversity in OBD scenario. Furthermore, we propose a novel yet simple algorithm, state density weighted (SDW) OBD, which emphasizes state diversity by weighting the distillation objective using the reciprocal of state density, thereby distilling a more diverse state information into synthetic data. Extensive experiments across multiple D4RL datasets confirm that SDW significantly enhances OBD performance when the original dataset exhibits limited state diversity.