robotdataset

robotdataset is a Python library for loading robot learning datasets from multiple sources (Google Cloud Storage, HuggingFace Hub) into the TorchRL TED format (TorchRL Episode Data), backed by memory-mapped tensors for efficient use in deep-learning training pipelines.

Documentation

Quick start

# All dataset backends (recommended)
pip install "robotdataset[all]"

# Or install only what you need:
pip install "robotdataset[oxe]"   # OXE / GCS datasets (requires TensorFlow)
pip install "robotdataset[hf]"    # HuggingFace datasets (Table30v2, AgiBotWorld-Beta)

from robotdataset import OXEDataset, TemporalSampler, batchViz

# Load two episodes of an OXE dataset (streams from GCS, caches locally)
dataset = OXEDataset(
    dataset_name="cmu_playing_with_food",
    episodes=[0, 2],
    batch_size=6,
)

# Attach a temporal sampler: each sample carries a 10-frame image history
sampler = TemporalSampler(
    delta_timestamps={"observation/image": [-0.9, -0.8, -0.7, -0.6, -0.5,
                                            -0.4, -0.3, -0.2, -0.1, 0.0]},
    control_frequency=10,
)
dataset.set_sampler(sampler)

batch = dataset.sample()
batch["observation/image"].shape   # (6, 10, 480, 640, 3) — (B, T, H, W, C)

# Visualize the batch as an animated mosaic
batchViz(batch, key="observation/image", fps=8)

Dataset status

All loaders are in alpha — APIs may change without notice between releases.

Public API

Everything below is importable directly from the top-level package:

from robotdataset import (
    # Dataset classes
    OXEDataset, OXEJAXDataset, Table30v2Dataset, AgiBotWorldBetaDataset, LeRobotVideoDataset,
    # Samplers
    TemporalSampler, EpisodeTubeletSampler, JAXTemporalSampler,
    # Discovery helpers
    list_datasets, validate_dataset_name, dataset2path, list_agibot_tasks,
    # Visualization
    batchViz, itemViz, episodeViz,
)

OXEJAXDataset and JAXTemporalSampler are None when JAX is not installed — they are optional-dependency imports.

License

MIT — see LICENSE.