简介

按 episode 把一份 LeRobot v3.0 数据集切成两份互不重叠的独立数据集（默认 train + val），源数据集不动。包 lerobot datasets/dataset_tools.py:144 split_dataset，外面加 crash-safe 与 skip-if-done。

功能描述

• 两种切法：行里给了 train_episode_indices 就按显式索引切（第二份拿剩下的全部 episode）；给 null 则按构造参数 train_fraction 切
• 比例模式的分配口径来自 lerobot _fractions_to_episode_indices：第一份取 int(总 episode 数 × fraction) 个，最后一份吃掉余数；episode 按索引顺序连续分配，不做 shuffle
• 每份 split 落在 <output_path>/<split_name>，是独立完整的 v3.0 数据集，episode 重索引成 0..N-1
• 重索引与元数据重写和 Episode 删除同一套逻辑：index / task_index 重编、任务表按该 split 里存活的任务重建、meta/stats.json 按该 split 的 episode 重新聚合、视频整文件保留时字节拷贝、混合文件用 PyAV 按帧区间重编码
• 两份 split 的逻辑数据集名由 lerobot 写成 <源 repo_id>_<split_name>
• crash-safe：开跑前把 output_path 整体删掉重建，中途失败再清一次半成品
• 完成后逐 split 校验都是合法 v3.0，缺任何一份直接报错
• skip-if-done：两份 split 都已是 v3.0 且 force=False 时返回 SKIPPED
• 校验：train_fraction 与 train_episode_indices 必须且只能生效一个（UDF 层保证行里给了索引就不传 fraction）；train_fraction 必须落在开区间 (0, 1)；索引越界、两份 split 有重叠、某份 split 为空都直接失败

算子参数

输入

输出

参数

注意事项

• output_path 是两份 split 的父目录，开跑前会被整体删除重建，不要指向源数据集或还有其他内容的目录。
• 比例是向下取整：4 个 episode + train_fraction=0.8 → train 3 / val 1。episode 太少时第一份可能被算成 0 个，lerobot 会跳过这份 split，随后本算子的完整性校验报错（体现为 Failed）。
• 划分是 episode 级、按索引顺序的连续切分，没有随机打乱；要随机划分请自己在行里给 train_episode_indices。
• 索引模式下第二份 split 拿剩余全部 episode；把所有 episode 都给了第一份会让第二份为空，上游直接报错。
• 每份 split 的 meta/info.json 里 splits 字段被 lerobot 统一写成 {"train": "0:N"}（N 是这份 split 自己的 episode 数），与 train/val 划分无关，不要拿它判断 split 身份。

调用示例

1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.embodied.lerobot.split_dataset_episodes_udf import (
10    SplitDatasetEpisodes,
11)
12
13if __name__ == "__main__":
14    if os.getenv("DAFT_RUNNER", "native") == "ray":
15        import ray
16        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
17        daft.set_runner_ray()
18    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
19
20# 夹具可用 test/_helpers/lerobot_edit_fixture.py:build_tabular_dataset 生成
21    # train_episode_indices 给 [0, 2] 即切到索引模式，train_fraction 失效
22    samples = {
23        "dataset_path": ["/path/to/lerobot_dataset"],
24        "output_path": ["/tmp/lerobot_split_out"],
25        "train_episode_indices": [None],
26    }
27    ds = daft.from_pydict(samples)
28    ds = ds.with_column(
29        "result",
30        aihc_udf(
31            SplitDatasetEpisodes,
32            construct_args={
33                "train_fraction": 0.5,
34                "split_names": ("train", "val"),
35                "force": False,
36            },
37            num_cpus=1,
38            concurrency=1,
39            batch_size=1,
40        )(col("dataset_path"), col("output_path"), col("train_episode_indices")),
41    )
42    ds.show()