Subtask annotation

Use subtask_annotation to add temporal subtask segments to robotics episodes, then use subtask_labeling to re-label those fixed segments with stronger previous/current/next visual context. The annotator samples an episode video into timestamped contact sheets, sends those images to a vision-language model, and writes the returned segments back to each row.

Subtask labels are useful when training policies that need a denser signal than the episode-level task description. They can identify coarse manipulation events such as reaching, grasping, moving, opening, pouring, or placing objects without manual per-episode labeling.

Run the following reference pipeline on your own infrastructure:

pipeline = (
    mdr.read_lerobot("hf://datasets/acme/demos")
    .map_async(
        mdr.robotics.subtask_annotation(
            video_key="observation.images.top",
            output_column="predicted_subtasks",
        ),
        max_in_flight=256,
    )
    .map_async(
        mdr.robotics.subtask_labeling(
            video_key="observation.images.top",
            segments_column="predicted_subtasks",
            output_column="labeled_subtasks",
        ),
        max_in_flight=256,
    )
)

By default, subtask_annotation uses Gemini through GoogleEndpointProvider. Set GOOGLE_GENERATIVE_AI_API_KEY before running the pipeline, or pass a different provider explicitly.

The labeling block is optional, but it is the recommended default when you want the best labels for fixed segments. If a segment already has a non-empty subtask, subtask_labeling will improve the labels using better conditioning. If segments only have timestamps, it runs a plain labeling prompt from previous/current/next visual context.

Input rows must implement RoboticsRow and expose the selected video through row.videos. If you start from another reader, convert rows with to_robot_rows(...) before annotation:

pipeline = (
    mdr.read_parquet("/data/episodes.parquet")
    .to_robot_rows(
        episode_id_key="episode_id",
        task_key="tasks",
        video_keys={"observation.images.top": "video_uri"},
        fps=30.0,
    )
    .map_async(
        mdr.robotics.subtask_annotation(
            video_key="observation.images.top",
        ),
        max_in_flight=256,
    )
)

The segmentation output column contains a list of segment dictionaries. Each segment has a start time, an end time, and a short action description under subtask:

[
    {"start_sec": 0.0, "end_sec": 1.2, "subtask": "reach object"},
    {"start_sec": 1.2, "end_sec": 2.8, "subtask": "grasp object"},
]

subtask_labeling writes fixed segments with a subtask value rewritten from a dedicated labeling pass. If the input segment has a non-empty subtask, the labeling prompt treats it as the seed label. If the input segment omits subtask or sets it to an empty string, the block uses a plain labeling prompt.

pipeline = pipeline.map_async(
    mdr.robotics.subtask_labeling(
        video_key="observation.images.top",
        segments_column="segments",
        output_column="labeled_subtasks",
    ),
    max_in_flight=256,
)

Video input is sent as contact sheets: timestamped image grids that preserve temporal context without sending the full video. This keeps requests smaller while giving the model enough visual context to choose event boundaries.

The default settings sample one frame every 0.5 seconds, resize each tile to 224px wide, and pack up to 20 frames per sheet in 5 columns.

For lower-level inference controls, see Generate Text and Multimodal and Structured Output.

For a complete local example, see Video Subtask Annotations.

For background and design decisions, see Annotating Robot Video Subtasks.