创建 LeRobot 2.1 数据集 • Axi's Blog

前言#

因为对于 LeRobot 数据集最近包括了大量的使用需求，但是寻找许久之后都没有非常清晰的文档或者教程，在这里给出最简单的教程，来解释如何创建 LeRobot 2.1 数据集。

安装#

首先，我们需要安装 LeRobot 2.1 的库，这里我们使用 pip 来安装：

pip install "lerobot @ git+https://github.com/huggingface/lerobot.git@2b71789e15c35418b1ccecbceb81f4a598bfd883"

bash

同时假如没有下载过，需要安装 ffmpeg

sudo apt update
sudo apt install ffmpeg

bash

创建数据集#

LeRobot 数据集一共包括两个环节，也不难理解，分别是创建数据集以及将每一个数据都存进去。

from lerobot.common.datasets.lerobot_dataset import HF_LEROBOT_HOME, LeRobotDataset

def create_dataset(
    repo_id: str,
    robot_type: str,
    mode: Literal["video", "image"] = "video",
) -> LeRobotDataset:
    motors = [
        "joint_0",
        "joint_1",
        "joint_2",
        "joint_3",
        "joint_4",
        "joint_5",
        "joint_6",
    ]
    cameras = [
        "camera_0",
        "camera_1",
        "camera_2",
    ]

    features = {
        "state.joints": {
            "dtype": "float32",
            "shape": (len(motors),),
            "names": [
                motors,
            ],
        },
        "action.joints": {
            "dtype": "float32",
            "shape": (len(motors),),
            "names": [
                motors,
            ],
        },
    }
    for cam in cameras:
        features[f"video.{cam}_view"] = {
            "dtype": mode,
            "shape": (3, 480, 640),  # (channels, height, width)
            "names": [
                "channels",
                "height",
                "width",
            ],
        }

    if Path(HF_LEROBOT_HOME / repo_id).exists():
        shutil.rmtree(HF_LEROBOT_HOME / repo_id)

    return LeRobotDataset.create(
        repo_id=repo_id,
        fps=15,
        robot_type=robot_type,
        features=features,
        use_videos=True,
        tolerance_s=0.0001,
        image_writer_processes=10,
        image_writer_threads=5,
        video_backend="ffmpeg",
    )

python

在这里我们使用将图像保存为视频，体现在输入的传参中，将 mode 设置为 "video"。如上的代码中定义了 joint 以及 camera 的基础信息。

本身对于 LeRobot 数据集的定义，主要是通过 features 来定义的，在此之后，features 就会被定义为一个 Dict，并且直接将其中存储数据即可。

与此同时，需要定义环境变量 export HF_LEROBOT_HOME=/path/to/your/lerobot/home，来设置 LeRobot 数据集的本地存储路径。

之后存储数据：

def process_single_dataset(
    dataset: LeRobotDataset,
    state_joints: np.ndarray,
    action_joints: np.ndarray,
    video_dict: dict[str, np.ndarray],
    instruction: str,
) -> LeRobotDataset:
    num_frames = state_joints.shape[0]

    for i in range(num_frames):
        frame = {
            "state.joints": state_joints[i],
            "action.joints": action_joints[i],
        }
        for camera, img_array in video_dict.items():
            frame[f"video.{camera}_view"] = img_array[i]
        dataset.add_frame(frame, task=instruction)

    dataset.save_episode()

    return dataset

python

在这里，只要使用同一个 dataset，LeRobot 会自动管理诸如 id 以及其他的统计信息。

小结#

LeRobot 数据集总体来说依然是比较好用的，在转录数据的时候，只需要从 Features 的角度入手即可，同时支持对于 Image 自动进行 Video 的编码，整体来说还是非常方便的。