Diversity You Can Actually Measure: A Fast, Model-Free Diversity Metric for Robotics Datasets

1. Problem Statement

Robotics datasets for imitation learning (IL) typically consist of long-horizon trajectories with varying lengths, comprising multiple modalities (states, actions, and high-dimensional observations like RGB video).

• The Challenge: Quantifying the diversity of such datasets is non-trivial. Standard pointwise distance metrics fail because they ignore the underlying trajectory structure and geometry. Furthermore, existing data curation methods often rely on training auxiliary models (e.g., VAEs, classifiers) or require access to policy rollouts, making them computationally expensive and unsuitable for large-scale pre-training datasets.
• The Goal: The authors aim to establish a principled, model-free metric to measure dataset diversity that respects trajectory structure. They seek to use this metric to curate subsets of demonstrations that maximize generalization performance in robot imitation learning without requiring additional policy training or rollout data.

2. Methodology

The core of the methodology involves extending information-theoretic entropy concepts to the space of trajectories using Signature Kernels.

A. Signature Transform and Kernel

• Signature Transform: The authors treat each demonstration as a path $x(t)$ in a high-dimensional space (combining states, actions, and embedded observations). The signature transform maps this path into a countably infinite sequence of iterated integrals. Crucially, this representation is invariant to reparameterization (i.e., speed variations do not change the signature), making it ideal for comparing trajectories of different lengths or speeds.
• Signature Kernel: A kernel function $k_{sig}(x, y)$ is defined as the inner product of the signatures of two paths. This provides a similarity measure that respects the geometric and sequential structure of the trajectories.

B. Defining Diversity via Entropy

The authors define diversity metrics based on the eigenvalues of the normalized Gram matrix (kernel matrix) $K_{sig}$ constructed from the dataset:

1. Signature Shannon Entropy:
   $$H_{sig}^{shannon} = -\sum_{\lambda_i \in \Lambda} \lambda_i \ln \lambda_i$$
   where $\Lambda$ are the eigenvalues of $K_{sig}/n$.
2. Signature von Neumann Entropy:
   $$H_{sig}^{von neumann} = -\text{Tr}\left(\frac{K_{sig}}{n} \ln \frac{K_{sig}}{n}\right)$$
   The paper proves that for this specific kernel construction, Shannon and von Neumann entropies are equivalent.
3. Determinant-based Diversity: The volume spanned by the samples is measured by $\det(K_{sig})$. A higher determinant implies the samples span a larger volume in the feature space (less redundancy).

C. The FAKTUAL Algorithm

FAKTUAL (FAst trajectory Kernel enTropy cUration for imitation Learning) is a data curation strategy designed to select a subset of size $m$ that maximizes these diversity metrics.

• Approach: Since finding the exact subset that maximizes entropy or determinant is an NP-hard combinatorial problem (Maximum Coverage), FAKTUAL employs a greedy algorithm (with optional local search) to find an approximate solution.
• Hybrid Strategy: To balance uniform randomness (entropy) with the inclusion of extreme/outlier points (volume), the algorithm uses a two-step selection:
  1. Select a subset of size $mp$ that maximizes Signature Entropy.
  2. From the remaining data, select a subset of size $m(1-p)$ that maximizes the Signature Determinant.
• Model-Free Nature: The algorithm requires no access to the imitation policy, no policy rollouts, and no auxiliary models. It operates solely on the raw dataset.

3. Key Contributions

1. Theoretical Extension: The authors extend Shannon and von Neumann entropy to robotics datasets by defining them over the eigenvalues of the Signature Kernel Gram matrix. This provides a mathematically principled way to measure the "information content" of trajectory data.
2. FAKTUAL Algorithm: They introduce a fast, model-free data curation algorithm that selects diverse subsets of demonstrations. It is computationally efficient, adding negligible overhead compared to policy training.
3. Empirical Validation: The paper demonstrates that diversity-aware curation correlates strongly with downstream policy success rates. FAKTUAL consistently outperforms random selection and matches or exceeds more complex, model-based curation methods in specific regimes.
4. Practical Tool: The work provides a practical tool for understanding dataset composition, showing that high entropy (diversity) is a strong predictor of policy performance, especially when data quality is already high.

4. Experimental Results

The authors evaluated FAKTUAL on:

• Simulated Benchmarks: RoboMimic (Can, Square, Transport) and MetaWorld (Door Open, Stick Push, Shelf Place).
• Real-World Tasks: Four manipulation tasks using an SO-ARM101 robot (Drawer Open, Mug Drag, Pick-and-Place Marker, Organize Objects).

Key Findings:

• Correlation with Success: There is a strong positive correlation (Pearson $r \approx 0.74 - 1.00$) between signature entropy and policy success rate.
• Performance vs. Random: FAKTUAL consistently outperforms random selection. In tasks like "Can" (RoboMimic), it achieves high success rates with significantly fewer demonstrations.
• Performance vs. State-of-the-Art:
  • FAKTUAL is generally better than random pruning.
  • It performs comparably to or slightly below complex, quality-based methods (like CUPID or Demo-SCORE) that require training auxiliary models or access to rollouts.
  • Crucial Advantage: FAKTUAL is significantly more computationally efficient. It does not require the expensive upfront cost of training auxiliary models, making it suitable for large-scale, multi-task pre-training datasets where other methods are too costly.
• Task Dependency: On simpler tasks, diversity curation yields immediate gains. On harder tasks, the benefit of diversity becomes more pronounced as the dataset size increases (once the "density" of successful behaviors is established).

5. Significance and Limitations

Significance:

• Scalability: FAKTUAL enables the curation of massive robotics datasets (e.g., Open X-Embodiment scale) without the computational bottleneck of training auxiliary scoring models.
• Theoretical Insight: It establishes a direct link between the geometric diversity of demonstration trajectories (measured via signature entropy) and the generalization capability of imitation learning policies.
• Complementarity: The authors argue that diversity-based curation (FAKTUAL) and quality-based curation (filtering bad data) are complementary. FAKTUAL can be used to ensure uniform coverage after high-quality data has been selected.

Limitations:

• Quality vs. Diversity: FAKTUAL does not inherently filter out low-quality or failed demonstrations. If a dataset contains many "bad" but diverse trajectories, FAKTUAL might select them. It assumes the input dataset is already of reasonable quality.
• Model Capacity: The benefits of diversity depend on the capacity of the policy model. If the model is too weak to capture the multimodality of the data, increasing diversity may not help or could even hurt performance.
• Adversarial Data: The method assumes the dataset does not contain adversarial or pathological trajectories that are radically different but harmful to learning.

In conclusion, the paper presents FAKTUAL as a practical, efficient, and theoretically grounded solution for curating diverse robotics datasets, proving that maximizing trajectory entropy is a viable proxy for improving imitation learning performance.