Skip to content

Load a LeRobot Dataset

SUMMARY

Use LeRobotDataset to open an existing LeRobot dataset from a local path or the Hugging Face Hub.

Loading can be customized with LeRobotDatasetConfig to select episodes, apply temporal sampling, configure transforms, and control decoding behavior.

Learn about LeRobot datasets
Review the dataset structure, feature schema, storage layout, and workflow.
View overview →

Method

Create a LeRobotDataset with a repository ID and an optional local path:

python
from telekinesis.dataengine import datasets

dataset = datasets.LeRobotDataset(
    repo_id="lerobot/pusht",
    local_path="results/lerobot/pusht",
)

If the dataset is not already available locally, it is resolved from the configured repository source and stored at the local path.

Once loaded, the dataset can be indexed directly:

python
frame = dataset[0]

Each item is returned as a frame dictionary containing the features defined by the dataset schema.

Parameter Configuration

ParameterTypeDefaultDescription
repo_idstrrequiredDataset repository identifier.
local_pathstr | Path | NoneNoneLocal directory used to load or store the dataset.
configLeRobotDatasetConfig | NoneNoneLoading, filtering, decoding, and transform configuration.

Config Reference

Use LeRobotDatasetConfig to control how the dataset is read:

python
config = datasets.LeRobotDatasetConfig(
    episode_indices=[0, 1, 2],
)

dataset = datasets.LeRobotDataset(
    repo_id="lerobot/pusht",
    local_path="results/lerobot/pusht",
    config=config,
)

For example, episode_indices limits the loaded dataset to a subset of episodes.

Other configuration options control temporal windows, transforms, video decoding, depth units, and cache behavior. For full confifuration details see Configuration.

Attribute Reference

AttributeTypeDescription
repo_idstrRepository identifier used to load the dataset.
local_pathPathResolved local directory containing the dataset.
episode_indiceslist[int] | NoneEpisode indices selected by the loading configuration. None means all episodes.
num_episodesintNumber of loaded or selected episodes.
num_framesintNumber of frames across the loaded or selected episodes.
featuresdict[str, dict]Feature schema describing the observations, actions, and other dataset values.

These can be used to verify which data was loaded before training, visualization, or further processing.

Indexing Reference

InputReturn TypeDescription
dataset[index]dictReturns one decoded frame.
dataset[start:stop]list[dict]Returns a list of decoded frames.

Example

python
"""Example script demonstrating how to load a LeRobot dataset using the Telekinesis Data Engine."""

from pathlib import Path

from loguru import logger

from telekinesis.dataengine import datasets


def load_lerobot_dataset_example():
    """Load selected episodes from a LeRobot dataset."""

    # 1. Define the dataset identity and local storage path.
    repo_id = "lerobot/pusht"

    local_path = (
        Path(__file__).resolve().parent.parent.parent.parent
        / "results"
        / repo_id
    )

    # 2. Configure how the dataset should be loaded.
    config = datasets.LeRobotDatasetConfig(
        episode_indices=[0, 1, 2],
    )

    # 3. Load the dataset.
    dataset = datasets.LeRobotDataset(
        repo_id=repo_id,
        local_path=local_path,
        config=config,
    )

    logger.info("LeRobot dataset loaded successfully.")
    logger.info(dataset)

    logger.info("Selected episodes: {}", dataset.episode_indices)
    logger.info("Number of selected episodes: {}", dataset.num_episodes)
    logger.info("Number of selected frames: {}", dataset.num_frames)

    # 4. Access the first frame.
    first_frame = dataset[0]

    logger.info("First frame: {}", first_frame)


if __name__ == "__main__":
    load_lerobot_dataset_example()

Next Steps