简介

给 meta/stats.json 里没有分位数（q01 / q10 / q50 / q90 / q99）的旧 v3.0 数据集补算统计。lerobot 只把这件事做成了 CLI 脚本，本算子把它包成库调用：判重用 has_quantile_stats，计算复用 compute_quantile_stats_for_dataset（逐 episode 遍历 dataset[idx] → get_feature_stats → aggregate_stats），结果用 lerobot write_stats 落盘。

源码锚点（lerobot 0.4.4）：datasets/v30/augment_dataset_quantile_stats.py:53 has_quantile_stats、:129 compute_quantile_stats_for_dataset、datasets/utils.py:295 write_stats。

功能描述

• 判重口径来自上游 has_quantile_stats：meta/stats.json 里任意一个特征带任意一个 q** 键就算「已有分位数」，overwrite=False 时直接返回 SKIPPED，并把现有 stats 原样回填到 features
• 计算走 lerobot 自己的 episode 循环：按 dataset_from_index / dataset_to_index 逐帧取 dataset[idx]，所以视频特征走的是数据集标准随机访问解码路径（真解码，不抽帧）
• 并发策略沿用上游：数据集不含 video key 时用最多 16 个线程并发处理 episode；含 video key 时强制顺序处理（视频解码非线程安全）
• 统计口径：跳过 string 特征；图像 / 视频的 uint8 数据先除 255 再在 (0, 2, 3) 上逐通道归约；分位数集合是 lerobot DEFAULT_QUANTILES = [0.01, 0.10, 0.50, 0.90, 0.99]
• 整份统计重算后覆盖写 meta/stats.json（不是只往里补 q* 键），mean / std / min / max 也一起刷新
• 与上游脚本的差异：不做 push_to_hub 与 hub 重打 tag（AIHC 数据集在 BOS 上，不进 HF hub）；不接受 repo_id，逻辑数据集名固定取数据集目录名
• 每个特征的统计字典以 JSON 字符串返回，输出 schema 不随数据集特征集变化

算子参数

输入

输出

参数

注意事项

• 分位数的聚合语义是「各 episode 分位数按帧数加权平均」，不是全量数据的真分位数；episode 内部又是 5000 bin 直方图近似。小 episode（几十帧量级）上与 np.quantile 有 1% 量级偏差，这是上游实现的口径，不是误差 bug。
• 内存：上游 process_single_episode 会把一个 episode 的所有帧 torch.stack 成一个张量，图像 / 视频数据集按「单 episode 全解开」估内存；本算子未改这个行为，长 episode 需要预留内存。
• 与「LeRobot 数据集统计量计算」的分工：本算子照抄 lerobot 的全量遍历口径（视频真解码、逐帧精确），代价是慢；只想快速刷统计用抽帧版算子。
• overwrite=False 的判重只看有没有 q* 键，不校验统计是否与当前数据一致；数据改过但 q* 键还在时不会重算，必须显式 overwrite=True。

调用示例

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.augment_quantile_stats_udf import (
10    AugmentQuantileStats,
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# 夹具可用 lerobot writer API 自建（见 test/cases/0.5.3/test_augment_quantile_stats.py）
21    samples = {
22        "dataset_path": ["/tmp/lerobot_quantile_out"],
23    }
24    ds = daft.from_pydict(samples)
25    ds = ds.with_column(
26        "result",
27        aihc_udf(
28            AugmentQuantileStats,
29            construct_args={"overwrite": False},
30            num_cpus=1,
31            concurrency=1,
32            batch_size=1,
33        )(col("dataset_path")),
34    )
35    ds.show()