Write to a LeRobot Dataset

SUMMARY

Use a writable LeRobotDataset to add feature-complete frames, save them as episodes, and finalize all pending data and metadata when writing is complete.

Learn about LeRobot datasets

Review the dataset structure, feature schema, storage layout, and workflow.

View overview →

Method

Add Frames

Call add_frame() once for each synchronized observation and action. The frame keys and values must match the feature schema defined when the dataset was created, and each frame must include task.

dataset.add_frame(frame)

frame = {
    "observation.images.camera1": camera_image,
    "observation.state": robot_state,
    "action": action,
    "task": "Pick and place the object",
}

dataset.add_frame(frame)

Save an Episode

After all frames for one task attempt have been added, persist the episode:

dataset.save_episode(parallel_encoding=True)

Saving writes the episode data, encodes visual streams, updates metadata, and clears the episode buffer for the next episode.

Discard Buffered Frames

Clear an incomplete or unwanted episode before it is saved:

dataset.clear_episode_buffer(delete_images=True)

Use has_pending_frames() to check whether the current episode buffer contains unsaved data.

Finalize the Dataset

Always finalize the dataset after the last episode, including when recording exits because of an error:

try:
    write_episodes(dataset)
finally:
    dataset.finalize()

WARNING

An unfinalized dataset can have invalid Parquet files because their footer metadata has not been written. A dataset cannot be indexed while recording; finalize it first.

Parameter Configuration

Function	Parameter	Type	Default	Description
add_frame	frame	dict	required	Feature-complete frame with a task entry. Torch tensors are converted to NumPy.
save_episode	episode_data	dict | None	None	Optional pre-built episode; otherwise uses buffered frames.
save_episode	parallel_encoding	bool	True	Encode multiple camera streams in parallel.
clear_episode_buffer	delete_images	bool	True	Also delete temporary images for the discarded episode.

Attribute Reference

Attribute	Return type	Description
has_pending_frames()	bool	Whether the current episode buffer contains unsaved frames.

Example

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

from pathlib import Path
import shutil

import numpy as np
from loguru import logger

from telekinesis.dataengine import datasets


def write_lerobot_dataset_example():
    """Programmatically create and write episodes to a LeRobot dataset."""

    # 1. Define the dataset identity, local storage path, and features.
    repo_id = "user/my_example_dataset"
    local_path = (
        Path(__file__).resolve().parent.parent.parent.parent
        / "results"
        / repo_id
    )
    # Remove the existing dataset directory to ensure rerun of example
    if local_path.exists():
        shutil.rmtree(local_path)

    features = {
        "observation.images.camera1": {
            "dtype": "video",
            "shape": [64, 64, 3],
            "names": ["height", "width", "channels"],
        },
        "observation.state": {
            "dtype": "float32",
            "shape": [6],
            "names": [
                "shoulder_pan_joint.pos",
                "shoulder_lift_joint.pos",
                "elbow_joint.pos",
                "wrist_1_joint.pos",
                "wrist_2_joint.pos",
                "wrist_3_joint.pos",
            ],
        },
        "action": {
            "dtype": "float32",
            "shape": [6],
            "names": [
                "shoulder_pan_joint.pos",
                "shoulder_lift_joint.pos",
                "elbow_joint.pos",
                "wrist_1_joint.pos",
                "wrist_2_joint.pos",
                "wrist_3_joint.pos",
            ],
        },
    }

    # 2. Create a writable LeRobot dataset.
    dataset = datasets.LeRobotDataset.create(
        repo_id=repo_id,
        local_path=local_path,
        fps=30,
        features=features,
        robot_type="ur10e",
        use_videos=True,
    )

    num_episodes = 3
    frames_per_episode = 5

    try:
        # 3. Programmatically add frames and save each episode.
        for episode_index in range(num_episodes):
            for _ in range(frames_per_episode):
                frame = {
                    "observation.images.camera1": np.random.randint(
                        0,
                        256,
                        size=(64, 64, 3),
                        dtype=np.uint8,
                    ),
                    "observation.state": np.random.rand(6).astype(np.float32),
                    "action": np.random.rand(6).astype(np.float32),
                    "task": "Dummy pick-and-place task",
                }

                dataset.add_frame(frame)

            dataset.save_episode()

            logger.info(f"Episode {episode_index+1} saved.")

    finally:
        # 4. Finalize all pending writers and metadata.
        dataset.finalize()

    logger.info("LeRobot dataset written successfully.")
    logger.info(dataset)
    logger.info(f"Local path: {local_path}")


if __name__ == "__main__":
    write_lerobot_dataset_example()

Next Steps

Load the Dataset

Reopen the finalized dataset and access its recorded frames.

Load dataset →

Inspect Dataset Metadata

Verify the saved episodes, features, statistics, and storage metadata.

Inspect metadata →

Visualize the Dataset

Inspect the recorded observations and actions in Rerun.

Visualize dataset →