Spaces:
Running on Zero
Running on Zero
File size: 7,261 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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """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
|