简介

闭集子任务标注算子，用本地 Qwen3-VL 对一段机器人操作视频做时序阶段划分，输出每个子任务的起止秒。移植 lerobot data_processing/sarm_annotations/subtask_annotation.py 的 VideoAnnotator（提示词来自同文件的 create_sarm_prompt），面向 SARM reward model 的训练数据标注。

与 AtomicActionSegment 的分工：本算子做语义 subtask 阶段划分（VLM 驱动），后者做运动学动作链切分（无模型启发式）。

功能描述

• 闭集词表由构造参数 subtask_list 给出并写进提示词，模型只能使用这些标签；提示词约束为：从 00:00 起无缝覆盖全片、相邻段 end 等于下一段 start、每个标签按时序恰出现一次、先写文字时间线再输出 JSON
• 推理前用 ffmpeg 把 [start_time_sec, end_time_sec] 区间重编码成 1 fps 的临时 mp4（libx264 / ultrafast / crf 23 / 去音轨），推理后删除临时文件；end_time_sec 为空时用 cv2 读帧数与帧率算出全片时长
• 生成参数沿用上游：max_new_tokens=1024、do_sample=True、temperature=0.7；回复先剥掉 markdown 代码块围栏，再 json.loads + pydantic 校验（解析不出来时用正则兜一次最外层花括号），失败按 max_retries 重试，仍失败抛 RuntimeError
• 模型类按 checkpoint 的 config 解析（AutoModelForImageTextToText，老版本 transformers 回落 AutoModelForVision2Seq）后注入 VideoAnnotator：上游写死了 Qwen3VLMoeForConditionalGeneration，换非 MoE 权重会在 from_pretrained 里对不上 config，所以这里自己加载再注入，model_name 可以指向同族任意尺寸的权重
• 权重两段式解析 model_path + model_name，构造算子时先检查 <model_path>/<model_name>/config.json 存在，缺权重直接 FileNotFoundError，不自动下载、不占 GPU
• 模型在首次标注时才加载（每个 worker 一份）；设备由 UDF 的 num_gpus 决定：>0 走 cuda，=0 走 cpu
• 模型输出的 MM:SS 时间戳按整数秒解析，duration_sec = end_sec - start_sec
• 行级容错：video_path 为空、视频不存在或标注失败都返回 Failed: ...，不打断整批
• 数据集级子任务进度先验 compute_temporal_proportions（SARM 论文公式 (1)）留在库层 daft.aihc.functions.embodied.lerobot.sarm_subtask_annotate，它要聚合多个 episode，不是逐行算子

算子参数

输入

输出

参数

注意事项

• 默认权重 Qwen3-VL-30B-A3B-Instruct 的 bf16 权重就有 58 GiB，24G 卡装不下。测试里按显存挑：≥64 GiB 用 30B MoE，≥20 GiB 用 Qwen/Qwen2.5-VL-7B-Instruct，更小用 Qwen/Qwen2.5-VL-3B-Instruct；两条路径走同一份算子代码。
• 权重不会自动下载，缺失时算子构造阶段直接抛 FileNotFoundError。
• 上游是采样解码（do_sample=True、temperature=0.7），同一段视频两次标注结果可能不同；边界是整数秒精度，亚秒级切分拿不到。
• 需要 ffmpeg 在 PATH 里，否则抽段阶段报 ffmpeg not found；抽段会往系统临时目录写一份 1 fps 的 mp4。
• result.raw_json 是解析并归一化之后的 subtasks，不是模型的原始回复文本。
• 库层 compute_temporal_proportions 沿用 lerobot 语义：单个 episode 内同名子任务的时长是覆盖而不是累加（同名只保留最后一段），最后按各名次均值归一化到和为 1。

调用示例

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.sarm_subtask_annotate_udf import SarmSubtaskAnnotate
10
11SUBTASK_LIST = [
12    "approach_object",
13    "grasp_object",
14    "lift_object",
15    "move_object",
16    "place_object",
17    "release_object",
18    "retreat",
19]
20
21if __name__ == "__main__":
22    if os.getenv("DAFT_RUNNER", "native") == "ray":
23        import ray
24        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
25        daft.set_runner_ray()
26    daft.set_execution_config(actor_udf_ready_timeout=3600, min_cpu_per_task=0)
27
28samples = {
29        "video_path": ["bos://your-bucket/sample.mp4"],
30        "start_time_sec": [0.0],
31        "end_time_sec": [40.0],
32    }
33    ds = daft.from_pydict(samples)
34    ds = ds.with_column(
35        "result",
36        aihc_udf(
37            SarmSubtaskAnnotate,
38            construct_args={
39                "subtask_list": SUBTASK_LIST,
40                "model_path": "/path/to/models",
41                "model_name": "Qwen/Qwen2.5-VL-7B-Instruct",
42                "dtype": "bfloat16",
43                "max_retries": 3,
44            },
45            num_cpus=4,
46            num_gpus=1,
47            concurrency=1,
48            batch_size=1,
49        )(col("video_path"), col("start_time_sec"), col("end_time_sec")),
50    )
51    ds.show()