so101-pick-place-tasks

# SO-ARM101 Pick-and-Place Manipulation Dataset ## 简介 本数据集是一组基于 [LeRobot v3.0](https://github.com/huggingface/lerobot) 格式构建的机器人操作数据集，使用 SO-ARM101 机械臂通过遥操作方式采集。数据集涵盖了多种日常物体的拾取、放置、归位以及方块堆叠等单臂操作任务，适用于模仿学习（Imitation Learning）、视觉-语言-动作（VLA）模型训练、策略学习等研究方向。 数据以 LeRobot v3.0 标准格式存储，包含高频状态记录、双目 RGB 视频流以及语言任务描述，可直接与 LeRobot 生态中的训练、评估和可视化工具链配合使用。仓库同时提供按任务合并后的 `merged_datasets/` 目录，每个子目录包含同一任务下所有采集日期的轨迹，便于直接用于策略训练。 ## 数据组织 仓库根目录下包含： - `merged_datasets/`：按任务合并后的数据集，每个子文件夹对应一个任务，文件夹名即任务描述（小写下划线格式）。 - 原始按日期划分的子数据集已合并到上述目录中。 - `check_lerobot_dataset.py` / `merge_lerobot_v3_datasets.py`：数据集检查与合并工具。 ## 数据集详情 | 属性 | 值 | |------|-----| | 机械臂平台 | SO-ARM101 | | 采集方式 | 遥操作（Teleoperation） | | 数据格式 | LeRobot v3.0 | | 视频帧率 | 30 FPS | | 图像分辨率 | 640 × 480（RGB） | | 动作维度 | 6（5-DOF 关节位置 + 1 夹爪开合） | | 轨迹总数 | 2,200 条 | | 总帧数 | 784,963 帧 | | 任务数量 | 9 个 | ### 合并后的子数据集 `merged_datasets/` 下包含 9 个按任务合并的数据集： | 子目录 | 任务描述 | 轨迹数 | 总帧数 | |--------|----------|--------|--------| | `pick_up_a_battery_and_place_in_the_bin` | Pick up a battery and place in the bin | 200 | 56,295 | | `pick_up_a_can_and_place_in_the_bin` | Pick up a can and place in the bin | 250 | 82,792 | | `pick_up_a_cube_and_place_in_the_bin` | Pick up a cube and place in the bin | 300 | 106,085 | | `pick_up_a_eraser_and_place_in_the_bin` | Pick up a eraser and place in the bin | 300 | 109,899 | | `pick_up_a_golf_and_place_in_the_bin` | Pick up a golf and place in the bin | 250 | 77,409 | | `pick_up_a_medicine_bottle_and_place_in_the_bin` | Pick up a medicine bottle and place in the bin | 300 | 116,342 | | `pick_up_a_plush_toy_and_place_in_the_bin` | Pick up a plush toy and place in the bin | 200 | 80,811 | | `Stack_the_cube_on_the_can` | Stack the cube on the can | 200 | 74,425 | | `Stack_the_smaller_cube_on_the_larger_one` | Stack the smaller cube on the larger one | 200 | 80,905 | ## 数据格式与结构 每个合并后的子数据集均遵循 LeRobot v3.0 目录结构： ``` merged_datasets/ └── pick_up_a_xxx_and_place_in_the_bin/ ├── meta/ │ ├── info.json # 数据集元信息、features 定义 │ ├── tasks.parquet # 任务索引与语言描述 │ ├── stats.json # 全局统计量 │ └── episodes/ │ └── episodes_*.parquet # episode 索引与长度 ├── data/ │ └── chunk-*/file-*.parquet # 逐帧数据（状态、动作、索引） └── videos/ ├── observation.images.top/ │ └── chunk-*/file-*.mp4 └── observation.images.wrist/ └── chunk-*/file-*.mp4 ``` ### Features | 字段 | 类型 | 形状 | 说明 | |------|------|------|------| | `action` | float32 | (6,) | 目标关节位置：shoulder_pan、shoulder_lift、elbow_flex、wrist_flex、wrist_roll、gripper | | `observation.state` | float32 | (6,) | 当前关节位置，维度与 action 一致 | | `observation.images.top` | video | (480, 640, 3) | 顶部俯视 RGB 视频流 | | `observation.images.wrist` | video | (480, 640, 3) | 腕部第一视角 RGB 视频流 | | `timestamp` | float32 | (1,) | 当前帧时间戳（秒） | | `frame_index` | int64 | (1,) | 当前 episode 内帧序号 | | `episode_index` | int64 | (1,) | 全局 episode 编号 | | `index` | int64 | (1,) | 全局帧编号 | | `task_index` | int64 | (1,) | 任务编号，对应 `tasks.parquet` | ## 使用方法 ### 环境准备 ```bash pip install lerobot ``` ### 使用 LeRobot 加载数据集 ```python from pathlib import Path from lerobot.common.datasets.lerobot_dataset import LeRobotDataset # 加载合并后的单任务数据集 dataset = LeRobotDataset.from_preloaded( repo_id="your-namespace/so101-pick-place-tasks", root=Path("./merged_datasets/pick_up_a_plush_toy_and_place_in_the_bin"), ) print(f"轨迹数: {dataset.num_episodes}") print(f"总帧数: {dataset.num_frames}") # 读取第 0 条轨迹 for frame in dataset.get_episode_buffer(0): print(frame["observation.images.top"].shape) # (480, 640, 3) print(frame["action"]) # (6,) break ``` ### 批量检查数据质量 仓库中提供了数据集完整性检查脚本： ```bash # 检查单个合并后的任务数据集 python check_lerobot_dataset.py --dataset ./merged_datasets/pick_up_a_plush_toy_and_place_in_the_bin # 批量检查 merged_datasets 下所有任务数据集 python check_lerobot_dataset.py --dataset ./merged_datasets --batch ``` ### 合并多个子数据集 仓库中提供了合并脚本，可将多个 LeRobot v3.0 子数据集合并为一个统一的数据集： ```bash python merge_lerobot_v3_datasets.py \ --input-dirs dataset_20260724_pick_place_eraser dataset_20260724_pick_place_plush_toy \ --output-dir merged_dataset ``` ## 数据收集与处理 ### 采集流程 1. **场景搭建**：在固定工作台上放置目标物体与收纳容器，顶部与腕部分别部署 RGB 相机。 2. **遥操作采集**：操作员通过示教设备实时控制 SO-ARM101 完成指定任务，同步记录关节状态、动作指令与双视角视频。 3. **数据切分**：以每次任务尝试为一条 episode，去除准备与结束阶段的无效帧后保存为 LeRobot v3.0 格式。 4. **质量校验**：使用 `check_lerobot_dataset.py` 检查时序连续性、视频帧数一致性、任务索引完整性等指标。 ### 标注信息 任务语言描述存储于 `meta/tasks.parquet`，通过 `task_index` 与数据帧关联。每个子数据集通常对应一个任务描述。 ## 许可协议 本数据集采用 [Apache-2.0](https://www.apache.org/licenses/LICENSE-2.0) 协议发布。 ## 引用 如果你在研究中使用了本数据集，请引用： ```bibtex @dataset{xarm_pick_place_2026, title={SO-ARM101 Pick-and-Place Manipulation Merged Dataset}, year={2026}, publisher={ModelScope}, note={LeRobot v3.0 format} } ``` ## 相关资源 - [LeRobot GitHub](https://github.com/huggingface/lerobot) - [LeRobot 数据集格式文档](https://github.com/huggingface/lerobot/blob/main/examples/12_convert_tsv_to_lerobot_dataset.py) - [ModelScope 数据集卡片规范](https://www.modelscope.cn/docs/datasets/dataset-card)