LeRobot documentation

将大型数据集迁移到 LeRobot 数据集 v3.0

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

将大型数据集迁移到 LeRobot 数据集 v3.0

本教程介绍如何将大规模机器人数据集迁移到 LeRobot Dataset v3.0 格式。我们将以 DROID 1.0.1 数据集作为主要示例,演示如何在 SLURM 集群上处理包含数千个分片的多 TB 级数据集。

文件组织:v2.1 与 v3.0

Dataset v3.0 从根本上改变了数据的组织和存储方式:

v2.1 结构(基于剧集):

dataset/
├── data/chunk-000/episode_000000.parquet
├── data/chunk-000/episode_000001.parquet
├── videos/chunk-000/camera/episode_000000.mp4
└── meta/episodes.jsonl

v3.0 结构(基于文件):

dataset/
├── data/chunk-000/file-000.parquet        # Multiple episodes per file
├── videos/camera/chunk-000/file-000.mp4   # Consolidated video chunks
└── meta/episodes/chunk-000/file-000.parquet  # Structured metadata

这种从单个剧集文件到基于文件的分块的转变,显著提升了性能并降低了存储开销。

Dataset v3.0 的新功能

Dataset v3.0 在处理大型数据集方面引入了重大改进:

🏗️ 增强的文件组织

  • 基于文件的结构:Episode 现在被分组到分块文件中,而不是单独的 episode 文件
  • 可配置的文件大小:适用于数据和视频文件
  • 提升存储效率:更好的压缩和更低的开销

📊 现代元数据管理

  • 基于 Parquet 的元数据:使用高效的 parquet 格式替换了 JSON Lines
  • 结构化剧集访问:通过 dataset.meta.episodes 直接访问 pandas DataFrame
  • 每回合统计:增强了回合级别的统计跟踪

🚀 性能增强

  • 内存映射访问:通过 PyArrow 内存映射改善 RAM 使用
  • 更快的加载速度:显著缩短了数据集初始化时间
  • 更好的可扩展性:专为包含数百万个 episode 的数据集设计

先决条件

在迁移大型数据集之前,请确保您已具备以下条件:

  • 已安装支持 v3.0 的 LeRobot。请遵循我们的 安装指南。
  • 足够的存储空间:原始数据集可能非常大(例如,DROID 需要 2TB)
  • 集群访问(推荐用于大型数据集):SLURM 或类似的作业调度器
  • 特定数据集的依赖项:对于 DROID,你需要 TensorFlow Dataset 实用工具

了解 DROID 数据集

DROID 1.0.1 是一个大规模机器人数据集的绝佳示例:

  • 大小:1.7TB(RLDS 格式),8.7TB(原始数据)
  • 结构:2048 个预定义的 TensorFlow 数据集分片
  • 内容:来自 Franka Emika Panda 机器人的 76,000+ 条机器人操作轨迹
  • 范围:跨多个环境和物体的真实世界操作任务
  • 格式:最初为 TensorFlow Records/RLDS 格式,需要转换为 LeRobot 格式
  • 托管:通过 gsutil 公开访问的 Google Cloud Storage

该数据集包含多种多样的操作演示,具体包括:

  • 多个摄像头视角(腕部摄像头、外部摄像头)
  • 自然语言任务描述
  • 机器人本体感知状态与动作
  • 成功/失败注释

DROID 功能模式

DROID_FEATURES = {
    # Episode markers
    "is_first": {"dtype": "bool", "shape": (1,)},
    "is_last": {"dtype": "bool", "shape": (1,)},
    "is_terminal": {"dtype": "bool", "shape": (1,)},

    # Language instructions
    "language_instruction": {"dtype": "string", "shape": (1,)},
    "language_instruction_2": {"dtype": "string", "shape": (1,)},
    "language_instruction_3": {"dtype": "string", "shape": (1,)},

    # Robot state
    "observation.state.gripper_position": {"dtype": "float32", "shape": (1,)},
    "observation.state.cartesian_position": {"dtype": "float32", "shape": (6,)},
    "observation.state.joint_position": {"dtype": "float32", "shape": (7,)},

    # Camera observations
    "observation.images.wrist_left": {"dtype": "video", "shape": (180, 320, 3), "names": ["height", "width", "channels"]},
    "observation.images.exterior_1_left": {"dtype": "video", "shape": (180, 320, 3), "names": ["height", "width", "channels"]},
    "observation.images.exterior_2_left": {"dtype": "video", "shape": (180, 320, 3), "names": ["height", "width", "channels"]},

    # Actions
    "action.gripper_position": {"dtype": "float32", "shape": (1,)},
    "action.cartesian_position": {"dtype": "float32", "shape": (6,)},
    "action.joint_position": {"dtype": "float32", "shape": (7,)},

    # Standard LeRobot format
    "observation.state": {"dtype": "float32", "shape": (8,)},  # joints + gripper
    "action": {"dtype": "float32", "shape": (8,)},  # joints + gripper
}

方法 1:单台计算机移植

步骤 1:安装依赖项

具体针对 DROID:

pip install tensorflow
pip install tensorflow_datasets

对于其他数据集,请为您的源格式安装相应的读取器。

步骤 2:下载原始数据

使用 gsutil 从 Google Cloud Storage 下载 DROID:

# Install Google Cloud SDK if not already installed
# https://cloud.google.com/sdk/docs/install

# Download the full RLDS dataset (1.7TB)
gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 /your/data/

# Or download just the 100-episode sample (2GB) for testing
gsutil -m cp -r gs://gresearch/robotics/droid_100 /your/data/

大型数据集需要大量的时间和存储空间:

  • 完整 DROID(1.7TB):根据带宽不同,下载可能需要数天
  • 处理时间:本地移植完整数据集需要 7 天以上
  • 上传时间:推送到 Hugging Face Hub 需要 3 天以上
  • 本地存储:处理后的 LeRobot 格式约需 400GB

步骤 3:移植数据集

python examples/port_datasets/port_droid.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1 \
    --push-to-hub

开发与测试

用于开发,你可以移植单个分片:

python examples/port_datasets/port_droid.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1_test \
    --num-shards 2048 \
    --shard-index 0

这种方法适用于较小的数据集或测试场景,但处理大型数据集则需要集群计算。

方法 2:SLURM 集群移植(推荐)

对于 DROID 这样的大型数据集,跨多个节点进行并行处理可以显著缩短处理时间。

步骤 1:安装集群依赖项

pip install datatrove  # Hugging Face's distributed processing library

步骤 2:配置您的 SLURM 环境

查找您的分区信息:

sinfo --format="%R"  # List available partitions
sinfo -N -p your_partition -h -o "%N cpus=%c mem=%m"  # Check resources

选择一个 CPU 分区 - 数据集移植无需 GPU。

步骤 3:启动并行移植任务

python examples/port_datasets/slurm_port_shards.py \
    --raw-dir /your/data/droid/1.0.1 \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
    --job-name port_droid \
    --partition your_partition \
    --workers 2048 \
    --cpus-per-task 8 \
    --mem-per-cpu 1950M

参数指南

  • --workers:并行任务数量(DROID 的分片数量上限为 2048)
  • --cpus-per-task:建议为帧编码并行化使用 8 个 CPU
  • --mem-per-cpu:加载原始帧所需的总内存约为 16GB(8×1950M)

先从较少的 worker 数量(例如 100 个)开始,以测试集群配置,然后再启动数千个任务。

步骤 4:监控进度

检查正在运行的任务:

squeue -u $USER

监控整体进度:

jobs_status /your/logs

检查单个作业日志:

less /your/logs/port_droid/slurm_jobs/JOB_ID_WORKER_ID.out

调试失败的作业:

failed_logs /your/logs/port_droid

步骤 5:聚合分片

所有移植任务完成后:

python examples/port_datasets/slurm_aggregate_shards.py \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
    --job-name aggr_droid \
    --partition your_partition \
    --workers 2048 \
    --cpus-per-task 8 \
    --mem-per-cpu 1950M

步骤 6:上传到 Hub

python examples/port_datasets/slurm_upload.py \
    --repo-id your_id/droid_1.0.1 \
    --logs-dir /your/logs \
    --job-name upload_droid \
    --partition your_partition \
    --workers 50 \
    --cpus-per-task 4 \
    --mem-per-cpu 1950M

Upload 使用更少的 worker(50 个),因为它受网络限制而非计算限制。

数据集 v3.0 文件结构

您完成的 dataset 将具有此现代结构:

dataset/
├── meta/
│   ├── episodes/
│   │   └── chunk-000/
│   │       └── file-000.parquet    # Episode metadata
│   ├── tasks.parquet               # Task definitions
│   ├── stats.json                  # Aggregated statistics
│   └── info.json                   # Dataset information
├── data/
│   └── chunk-000/
│       └── file-000.parquet        # Consolidated episode data
└── videos/
    └── camera_key/
        └── chunk-000/
            └── file-000.mp4        # Consolidated video files

这将旧的按文件划分剧集的结构替换为高效且大小最优的块。

从 Dataset v2.1 迁移

如果您拥有 v2.1 格式的现有数据集,请使用迁移工具:

python src/lerobot/scripts/convert_dataset_v21_to_v30.py \
    --repo-id your_id/existing_dataset

这将自动执行以下操作:

  • 将文件结构转换为 v3.0 格式
  • 将元数据从 JSON Lines 迁移到 parquet
  • 汇总统计数据并创建每回合的统计信息
  • 更新版本信息

性能优势

Dataset v3.0 为大型数据集提供了显著改进:

  • 更快的加载速度:初始化时间减少 3-5 倍
  • 内存效率:通过内存映射实现更优的 RAM 使用
  • 可扩展处理:高效处理数百万个回合
  • 存储优化:减少文件数量并改进压缩
Update on GitHub