LeRobot Dataset

SUMMARY

LeRobotDataset records teleoperated episodes - camera frames, proprioceptive state, and the action taken at each step - directly into the on-disk LeRobot dataset format (parquet + MP4, codebase v3.0), and reads them back as a PyTorch Dataset for training. It is the same dataset format Cerebellum's ACT, π0.5, and Fast-WAM policies fine-tune on.

WARNING

🚧 LeRobot dataset support is under active development.

The interface documented here reflects the intended API and is subject to change before general availability.

Install

pip install telekinesis-dataengine

Import

from telekinesis.dataengine import LeRobotDataset

---

Capabilities

Feature	Description
Record	create() starts a new dataset; add_frame() buffers a frame; save_episode() writes the buffered episode's parquet rows and encodes its camera videos.
Resume	resume() reopens an existing local or Hub dataset and continues appending episodes to it.
Read	Indexing (dataset[i]) returns a frame with images decoded and transforms applied, exactly like any torch.utils.data.Dataset.
Temporal windows	delta_timestamps pulls a window of past/future frames per feature (e.g. an action chunk) alongside the current one.
Depth cameras	Depth streams are encoded separately from RGB (12-bit HEVC by default) with a configurable meters/millimeters output unit.
Hub integration	push_to_hub() uploads a recorded dataset, including videos and a generated dataset card.
Dataset tools	Module-level helpers to merge, split, delete episodes from, re-encode, or recompute statistics for an existing dataset.

Quickstart

from telekinesis.dataengine import LeRobotDataset

# ------------------------------------------------
# 1. Define the feature schema and create the dataset
# ------------------------------------------------
features = {
    "observation.images.cam_high": {"dtype": "video", "shape": (480, 640, 3), "names": ["height", "width", "channel"]},
    "observation.state":           {"dtype": "float32", "shape": (14,), "names": None},
    "action":                       {"dtype": "float32", "shape": (14,), "names": None},
}

dataset = LeRobotDataset.create(
    repo_id="telekinesis-ai/insert_gear",
    fps=30,
    features=features,
    robot_type="aloha",
)

# ------------------------------------------------
# 2. Record an episode frame by frame
# ------------------------------------------------
for observation, action in teleoperated_episode:
    dataset.add_frame({
        "observation.images.cam_high": observation.cam_high,
        "observation.state": observation.state,
        "action": action,
        "task": "insert the gear into the housing",
    })

dataset.save_episode()

# ------------------------------------------------
# 3. Finalize when the recording session is done
# ------------------------------------------------
dataset.finalize()

# ------------------------------------------------
# 4. Read it back for training
# ------------------------------------------------
frame = dataset[0]
print(dataset.num_episodes, dataset.num_frames)