LeRobot Dataset

custom_dataset.py

from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
from lerobot.configs.train import TrainPipelineConfig
from lerobot.configs.default import DatasetConfig
from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
import numpy as np
import os
from pathlib import Path


def main():
    dataset_name = "example_dataset"
    this_dir = os.path.dirname(os.path.abspath(__file__))
    dataset_root = os.path.join(this_dir, dataset_name)

    features = {
            "action": {
                "dtype": "float32",
                "shape": (6,),
                "names": ["x", "y", "z", "roll", "pitch", "yaw"],
            },
            "observation.state": {
                "dtype": "float32",
                "shape": (6,),
                "names": ["x", "y", "z", "roll", "pitch", "yaw"],
            },
            "observation.images.camera1": {
                "shape": (480, 640, 3),
                "names": ["height", "width", "channels"],
                "info": None,
                "dtype": "image",
            },
        }

    dataset = None
    try:
        dataset = LeRobotDataset.create(dataset_name, 10, features=features, root=dataset_root)
    except FileExistsError:
        dataset = LeRobotDataset(dataset_name, root=dataset_root)

    for _ in range(2):
        dataset.add_frame(
            {
                "observation.images.camera1": np.zeros((480, 640, 3), dtype=np.uint8),
                "observation.state": np.zeros((6,), dtype=np.float32),
                "action": np.zeros((6,), dtype=np.float32),
                "task": "example_task",
            }
        )
    dataset.save_episode()

    config = TrainPipelineConfig(
        dataset=DatasetConfig(
            repo_id=dataset_name,
            root=dataset_root,
        ),
        policy=DiffusionConfig(),
    )
    config._save_pretrained(Path(dataset_root))


# python ~/lerobot/lerobot/scripts/visualize_dataset.py --repo-id example_dataset --mode local --root ${PWD}/example_dataset --episode-index 0
# python ~/lerobot/lerobot/scripts/train.py --config_path ./example_dataset/train_config.json
if __name__ == "__main__":
    main()