State Diversity Matters in Offline Behavior Distillation

TL;DR

The paper studies how dataset characteristics affect offline behavior distillation (OBD). It shows a misalignment where high-quality original data does not guarantee a superior distilled dataset, especially under large training loss, and argues that state diversity better supports learning in the underfitting regime. A theoretical framework partitions states into pivotal and surrounding categories and demonstrates that surrounding error can dominate policy performance when pivotal error is large. To exploit state diversity, the authors propose State Density Weighted (SDW) OBD, weighting the distillation loss by the state density, and demonstrate substantial gains on D4RL benchmarks, particularly for datasets with limited diversity.

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.

Figures (5)

• Figure 1: t-SNE visualization of M-R (blue) and M-E (red) states.
• Figure 2: Plots of policy performance (normalized return) as functions of the loss w.r.t. real data. Each point is calculated and then averaged over five trials.
• Figure 3: Policy gap bound as a function of (a) pivotal error $\epsilon$ under different surrounding error $\epsilon_\mu$; (b) distillation loss. $R$ denotes $R_{\max}$ for simplicity.
• Figure 4: The distribution of state density $d(s)$. The y-axis represents the density of the value $\log d(s)$, analogous to a probability density function.
• Figure 5: Plots of SDW performance (normalized return) as functions of $\tau$. Each point is averaged on five trials.

Theorems & Definitions (11)

• Definition 1: Pivotal and surrounding state
• Definition 2
• Definition 3
• Theorem 1: Theorem 2.1 in ross2010efficient
• Remark 1
• Theorem 2: Main Theorem
• proof
• Corollary 1
• proof
• Lemma 1