"""Per-dataset annotation loss specification. A schema can declare one or more `AnnotationLossSpec` entries to request auxiliary losses computed on annotation columns present in the parquet data. Dataset adapters pass through arbitrary string columns (e.g. RoboInter's `annotation.substask`); this spec tells the training pipeline which ones to tokenize and supervise with a text next-token CE loss. Only datasets whose parquet actually contains the named `field` should declare the spec. Datasets with no annotation_losses (e.g. oxe-auge) take the fast pure-MSE path — completely decoupled from this feature. Decoupling from Knowledge Isolation ----------------------------------- Annotation CE is orthogonal to the KI toggle: - KI=false: annotation CE flows to VLM (standard π0 joint-grad, same as MSE). - KI=true : annotation CE still flows to VLM (MSE detach does not affect it — KI only detaches the MSE path into DiT). The whole point of annotation CE is to *train the VLM on semantic labels*, so it must always be un-detached. Only MSE detaching is governed by KI. """ from __future__ import annotations from dataclasses import dataclass from typing import Any _SUPPORTED_LOSS_TYPES = ("ce_text",) # The raw per-field columns that BuildUnifiedAnnotationTransformFn packs into # the synthesized `annotation.unified` text. Both the transform (builder) and # the v3.0 adapter's column projection (which must fetch these physical columns # so the builder has inputs) consume it. # # Field order matters (it is the order the model sees): bbox/object-location # fields come BEFORE annotation.substask per π0.5 App B.2 "locate first, then # subtask". `annotation.segmentation` is intentionally OMITTED — a dense pixel # mask is too heavy to tokenize as text. NOTE: "substask" (double s) is the # dataset's REAL column spelling. UNIFIED_ANNOTATION_FIELDS: tuple[str, ...] = ( "annotation.object_box", "annotation.gripper_box", "annotation.affordance_box", "annotation.substask", "annotation.primitive_skill", "annotation.instruction_add", "annotation.placement_proposal", "annotation.state_affordance", "annotation.contact_frame", "annotation.contact_points", "annotation.trace", # annotation.time_clip kept available but lower-signal — included for # completeness, model will learn to skip if useless. "annotation.time_clip", ) @dataclass(frozen=True) class AnnotationLossSpec: """Frozen spec for a single annotation-driven auxiliary loss. Fields: field: Parquet column name, e.g. "annotation.substask". Must be a string-typed column; the adapter passes its raw value through to the sample dict unchanged. loss_type: Only "ce_text" in v1 — tokenize the string with the VLM tokenizer and apply next-token cross entropy on VLM hidden states at the annotation positions. weight: Multiplier applied to this CE term when added to the total loss. Must be > 0. Typical range [0.1, 1.0]. max_length: Max token count after tokenization. Shorter strings are padded; longer ones are truncated. Padding positions are excluded from the CE computation via an attention-like mask. """ field: str loss_type: str = "ce_text" weight: float = 0.5 max_length: int = 32 def __post_init__(self) -> None: if not self.field or not isinstance(self.field, str): raise ValueError( f"AnnotationLossSpec.field must be a non-empty string, got {self.field!r}" ) if self.loss_type not in _SUPPORTED_LOSS_TYPES: raise ValueError( f"AnnotationLossSpec.loss_type must be one of {_SUPPORTED_LOSS_TYPES}, " f"got {self.loss_type!r}" ) # bool is a subclass of int, so ``isinstance(True, (int, float))`` is # True. Without the explicit bool guard a config typo like # ``weight: true`` / ``max_length: false`` would be silently coerced to # 1.0 / 1 / 0 and change the loss semantics instead of being rejected. if (isinstance(self.weight, bool) or not isinstance(self.weight, (int, float)) or self.weight <= 0): raise ValueError( f"AnnotationLossSpec.weight must be a real number > 0, got " f"{type(self.weight).__name__}={self.weight!r}" ) if (isinstance(self.max_length, bool) or not isinstance(self.max_length, int) or self.max_length <= 0): raise ValueError( f"AnnotationLossSpec.max_length must be a positive int, got " f"{type(self.max_length).__name__}={self.max_length!r}" ) def to_dict(self) -> dict[str, Any]: return { "field": self.field, "loss_type": self.loss_type, "weight": float(self.weight), "max_length": int(self.max_length), } @classmethod def from_dict(cls, d: dict[str, Any]) -> "AnnotationLossSpec": # Validate the RAW types before float()/int() coercion. Otherwise # ``"weight": true`` round-trips through float(True)==1.0 (and # ``"max_length": false`` through int(False)==0) and silently changes # semantics instead of raising. bool is a subclass of int, so it is # excluded explicitly here as well as in __post_init__. weight_raw = d.get("weight", 0.5) if isinstance(weight_raw, bool) or not isinstance(weight_raw, (int, float)): raise ValueError( f"AnnotationLossSpec.weight must be a real number, got " f"{type(weight_raw).__name__}={weight_raw!r}" ) max_length_raw = d.get("max_length", 32) if isinstance(max_length_raw, bool) or not isinstance(max_length_raw, int): raise ValueError( f"AnnotationLossSpec.max_length must be an int, got " f"{type(max_length_raw).__name__}={max_length_raw!r}" ) return cls( field=d["field"], loss_type=d.get("loss_type", "ce_text"), weight=float(weight_raw), max_length=int(max_length_raw), ) def batch_key_tokens(self) -> str: """Sample-dict / batch-dict key for this annotation's token IDs.""" return f"annotation_tokens__{self.field}" def batch_key_mask(self) -> str: """Sample-dict / batch-dict key for this annotation's valid-token mask.""" return f"annotation_mask__{self.field}" def batch_key_weight(self) -> str: """Sample-dict / batch-dict key carrying `spec.weight` as a scalar tensor. Per-sample scalar lets the model compose loss without a schema reference plumbed through the training loop.""" return f"annotation_weight__{self.field}" @staticmethod def field_from_tokens_key(key: str) -> str | None: """Reverse of `batch_key_tokens` — extract field name or None.""" prefix = "annotation_tokens__" return key[len(prefix):] if key.startswith(prefix) else None