Spaces:
Sleeping
Sleeping
File size: 5,267 Bytes
34f3bc9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """Per-frame subtask annotation builder for AgiBot World samples.
agibot_dual_arm has exactly one task description per task_NNN dir (1-row
tasks.parquet). The v30 adapter already resolves ``task_index`` → task string
into ``data["task"]``. Naively reusing that as the prediction target creates a
pathological loop: the prompt is a generic VLA
template that often substitutes ``task`` as the user-text instruction → the
model is told the answer in its input.
This transform:
1. Copies ``data["task"]`` into ``data["annotation.subtask"]`` (the field
declared by ``schemas/agibot_dual_arm.SCHEMA.annotation_losses``).
2. Rewrites ``data["task"]`` to a generic descriptor so the prompt asks the
model to predict the task instead of repeating it.
The transform is hydrated via ``hydrate_all`` based on schema id — for non-
agibot schemas it is a no-op.
Pipeline placement (set by ``LabVLADatasetConfig.__post_init__``):
... → unify_anno → BuildAgiBotSubtaskTransformFn → qwen_processor → ...
placing it AFTER unify_anno guarantees no collision with the 12-field RoboInter
unified-annotation builder, and BEFORE qwen_processor so ``data["task"]``
rewrite reaches the prompt builder.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass, replace
from src.transforms.core import DataDict, DataTransformFn
# Conservative generic prompt; short tokens to keep prefix budget tight.
DEFAULT_GENERIC_PROMPT = "Describe the high-level task being performed."
@DataTransformFn.register_subclass("build_agibot_subtask")
@dataclass
class BuildAgiBotSubtaskTransformFn(DataTransformFn):
"""Promote ``data['task']`` to ``data['annotation.subtask']`` for agibot.
Inactive by default; ``hydrate_all`` flips ``enabled=True`` for schemas
whose ``annotation_losses`` declare a spec with ``field='annotation.subtask'``.
"""
enabled: bool = False
generic_prompt: str = DEFAULT_GENERIC_PROMPT
rewrite_task: bool = True # set False to keep original task in prompt
@staticmethod
def _is_present(value) -> bool:
"""True if ``value`` is a real (non-empty, non-null-sentinel) string.
Mirrors the null-stringification artifacts handled elsewhere in the
annotation pipeline so a parquet "None"/"nan"/"[]" cell counts as
MISSING rather than a real existing subtask.
"""
if value is None:
return False
try:
s = str(value).strip()
except Exception:
return False
return bool(s) and s.lower() not in {"none", "nan", "null", "[]", "{}", "''", '""'}
def __call__(self, data: DataDict) -> DataDict:
if not self.enabled:
return data
# Only backfill when annotation.subtask is MISSING. A sample that
# already carries a real subtask annotation is authoritative —
# overwriting it with data["task"] and clobbering the prompt would
# corrupt schemas/samples that legitimately provide subtask.
if self._is_present(data.get("annotation.subtask")):
return data
raw = data.get("task", "")
try:
task_str = str(raw).strip() if raw is not None else ""
except Exception:
task_str = ""
if not task_str:
# Missing task — skip silently. AnnotationTokenizeTransformFn
# treats absent annotation field as "no CE for this sample".
return data
data["annotation.subtask"] = task_str
if self.rewrite_task:
data["task"] = self.generic_prompt
return data
def hydrate(self, ctx) -> "BuildAgiBotSubtaskTransformFn":
# Gate on an EXPLICIT schema identity, not just the presence of an
# `annotation.subtask` loss field: gating on the field alone would fire
# for ANY schema reusing that field name, clobbering its `task` with the
# generic prompt and overwriting its `annotation.subtask`. The
# task→subtask backfill + prompt rewrite is AgiBot-specific (one task
# description per task dir), so require the AgiBot schema identity.
# robot_type / schema_id carry the family tag.
schema = ctx.schema
_has_subtask_loss = any(
getattr(spec, "field", None) == "annotation.subtask"
for spec in (schema.annotation_losses or ())
)
_robot_type = str(getattr(schema, "robot_type", "") or "")
_schema_id = str(getattr(schema, "schema_id", "") or "")
_is_agibot_schema = (
_robot_type.startswith("agibot") or _schema_id.startswith("agibot")
)
_agibot_enabled = bool(_has_subtask_loss and _is_agibot_schema)
t = replace(self, enabled=_agibot_enabled)
if _has_subtask_loss and not _is_agibot_schema:
logging.info(
f"{t.__class__.__name__} left disabled: schema "
f"{_schema_id!r} declares annotation.subtask but is not an "
f"AgiBot schema — skipping task→subtask backfill/prompt rewrite."
)
if _agibot_enabled:
logging.info(
f"Hydrated {t.__class__.__name__} enabled=True ({schema.schema_id})"
)
return t
|