robotdataset 0.2.3

A TorchRL wrapper for robot learning datasets.

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

Page	Covers
Overview & TED format	Library scope, architecture, the TED format, caching
OXE datasets	OXEDataset, list_datasets, validate_dataset_name, dataset2path
OXE JAX datasets	OXEJAXDataset — NumPy/JAX batch path
Table30v2	Table30v2Dataset (RoboChallenge/Table30v2, HuggingFace)
AgiBotWorld-Beta	AgiBotWorldBetaDataset, list_agibot_tasks
Samplers	TemporalSampler, EpisodeTubeletSampler, JAXTemporalSampler
Visualization	batchViz, itemViz, episodeViz
Config	Config — fluent config builder, YAML/JSON loading, W&B sweep export

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.

Dataset	Class	Source	Status
Open X-Embodiment (OXE)	OXEDataset	gs://gresearch/robotics	Alpha (usable, not stable)
Open X-Embodiment (JAX path)	OXEJAXDataset	gs://gresearch/robotics	Alpha (JAX path)
Table30 v2	Table30v2Dataset	RoboChallenge/Table30v2 (HF)	Alpha
AgiBotWorld-Beta	AgiBotWorldBetaDataset	agibot-world/AgiBotWorld-Beta (HF)	Alpha
LIBERO	—	openvla/modified_libero_rlds (HF)	Planned

Public API

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

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

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

License

MIT — see LICENSE.