kite-data

Dataset loading & normalization for robotics ML.

Part of the kite-ml package family.

Install

pip install kite-data              # core (numpy + pyarrow)
pip install kite-data[huggingface] # + HuggingFace Hub support
pip install kite-data[s3]          # + S3 support
pip install kite-data[all]         # everything

Quick start

from kite_data import load_dataset, convert_dataset, inspect_dataset

# Load into memory (returns Iterator[Episode])
for episode in load_dataset("lerobot/pusht"):
    print(episode.observations["state"].shape)  # (T, 5)
    print(episode.actions.shape)                # (T, 2)

# Convert to canonical Parquet on disk
convert_dataset("lerobot/pusht", output="./my_data/")

# Inspect schema without downloading full data
info = inspect_dataset("lerobot/pusht")
print(info["num_episodes"])  # 206

# S3 with credentials
for episode in load_dataset("s3://my-bucket/demos/", aws_access_key_id="...", aws_secret_access_key="..."):
    ...

Canonical output format

All data is normalized to Apache Parquet with a sidecar directory for media:

dataset/
├── meta.json           # Schema + dataset-level stats
├── episodes.parquet    # One row per timestep
└── media/              # Images/videos (only if present)
    └── ep_000/
        ├── frame_000.png
        └── ...

Supported sources

Source	URI format
HuggingFace Hub	lerobot/pusht or https://huggingface.co/datasets/...
Amazon S3	s3://bucket/prefix/

Supported formats

Format	Auto-detected by
LeRobot	meta/info.json present
D4RL	.hdf5 with observations key
Raw/Custom	Fallback — user provides column map

License

Apache 2.0