这是一个用于合并 OpenPI 数据集的工具脚本。它可以将一个源数据集(目录 B)合并到目标数据集(目录 A)中,同时保持数据的完整性和一致性。

• 自动合并数据目录中的 episode 文件
• 更新并合并元数据文件(episodes.jsonl, episodes_stats.jsonl, tasks.jsonl, info.json)
• 自动处理索引偏移和重新编号
• 验证机器人类型一致性
• 支持任务去重(计算任务集合的并集)

python openpi_dataset_merge.py --dir_a "目标目录路径" --dir_b "源目录路径"

python openpi_dataset_merge.py \
  --dir_a "/path/to/target/dataset" \
  --dir_b "/path/to/source/dataset"

• --dir_a: 目标目录路径(数据将被合并到此目录)
• --dir_b: 源目录路径(数据将从此目录读取并合并)

两个目录都应该具有以下结构:

dataset/
├── data/
│   ├── chunk-000/
│   │   ├── episode_000000.parquet
│   │   ├── episode_000001.parquet
│   │   └── ...
│   └── chunk-001/
│       └── ...
└── meta/
    ├── episodes.jsonl
    ├── episodes_stats.jsonl
    ├── tasks.jsonl
    └── info.json

脚本将按照以下步骤执行:

1. 合并元数据目录: 合并 episodes.jsonl, episodes_stats.jsonl, tasks.jsonl 和 info.json
2. 修订统计信息: 更新 episodes_stats.jsonl 中的索引和任务索引
3. 合并数据目录: 将源目录的 episode 文件复制到目标目录,并更新相关索引

• 合并前请确保两个数据集的 robot_type 一致,否则会报错
• 建议在合并前备份目标目录
• 每个 chunk 目录最多包含 1000 个 episode 文件
• 合并过程会自动处理全局帧索引的偏移

• Python 3.x
• pandas
• numpy
• draccus

pip install pandas numpy draccus