"""Dataset Blueprint — declarative single-source-of-truth for authoring schemas. Users describe a dataset with high-level constructs (arm DoF, parquet columns, camera roles, annotations). `.build()` compiles these into a `DatasetSchema` — the same internal dataclass that every downstream module (stats, merge, preflight, adapter, transform chain, model, deploy) already consumes. Two top-level blueprints cover every dataset we have: SingleArmBlueprint — arm + gripper layout (oxe_auge, labutopia_level3, future droid-family datasets). MultiChannelBlueprint — arbitrary list of (state_col, action_col, dim, mode) channels (robocoin's action + dedicated gripper_open_scale_action). Anything a DatasetSchema needs but the blueprint can't infer (e.g. canonical dual-arm layout on a multi-channel robocoin) is passed explicitly as an argument. The blueprint compiler is a pure function — running it twice yields bit-identical DatasetSchema.to_dict() — which makes stats.json / labvla_manifest.json / checkpoint compatibility self-verifying. Normalization mode ------------------ The project-wide default is ``"mean_std"`` — z-score, set at the transform chain level in ``policies/LabVLA/configuration_labvla.py``. ``hydrate_all`` in ``transforms/core.py`` then auto-injects per-dim ``"q01_q99"`` overrides on the canonical gripper indices declared by ``schema.gripper_action_dims``, so the effective policy is "arm dims → mean_std, gripper dims → q01_q99". Schemas don't carry a per-dim normalization field; overrides flow through ``schema.gripper_action_dims`` (and ``arm_layout.gripper_indices_canonical``, which mirrors it). """ from __future__ import annotations from dataclasses import dataclass, field from typing import Mapping, Optional from .annotation_loss import AnnotationLossSpec from .arm_layout import ArmLayoutSpec, ArmCount from .camera_mapping import expand_camera_mapping from .dataset_schema import DatasetSchema # --------------------------------------------------------------------------- # # Authoring primitives # # --------------------------------------------------------------------------- # @dataclass(frozen=True) class AnnotationSpec: """Blueprint-level shorthand for an annotation CE loss declaration. Compiled to `schema.annotation_loss.AnnotationLossSpec` at build() time. """ field: str weight: float = 0.5 max_length: int = 32 loss_type: str = "ce_text" def to_loss_spec(self) -> AnnotationLossSpec: return AnnotationLossSpec( field=self.field, loss_type=self.loss_type, weight=self.weight, max_length=self.max_length, ) @dataclass(frozen=True) class SingleArm: """Single-arm spec: arm joints + one gripper channel. The four parquet columns fully describe the per-frame layout. Gripper state/action may live in the *same* column as arm joints (e.g. oxe_auge's `observation.joints` holds both) — in that case the user sets the same column name for both `joint_*` and `gripper_*`; the blueprint infers that only one state/action key is needed. """ dof: int joint_state_col: str joint_action_col: str gripper_state_col: str gripper_action_col: str arm_mode: str = "delta" # "delta" | "abs" gripper_mode: str = "abs" # "delta" | "abs" def __post_init__(self) -> None: if self.dof <= 0: raise ValueError(f"SingleArm.dof must be > 0, got {self.dof}") for name, mode in (("arm_mode", self.arm_mode), ("gripper_mode", self.gripper_mode)): if mode not in ("delta", "abs"): raise ValueError( f"SingleArm.{name} must be 'delta' or 'abs', got {mode!r}" ) @dataclass(frozen=True) class StateActionChannel: """One independent (state_col, action_col) channel with its own dim/mode. Used by MultiChannelBlueprint for datasets whose action vector is split across multiple parquet columns (robocoin: `action` for arm + `gripper_open_scale_action` for gripper, with different dims and modes). """ state_col: str action_col: str dim: int mode: str = "delta" # "delta" | "abs" is_gripper: bool = False # marks all `dim` flat indices as gripper dims def __post_init__(self) -> None: if self.dim <= 0: raise ValueError(f"StateActionChannel.dim must be > 0, got {self.dim}") if self.mode not in ("delta", "abs"): raise ValueError( f"StateActionChannel.mode must be 'delta' or 'abs', got {self.mode!r}" ) if self.is_gripper and self.mode == "delta": raise ValueError( f"is_gripper=True channel must use mode='abs' (gripper never delta). " f"Got mode={self.mode!r}, state_col={self.state_col!r}" ) # --------------------------------------------------------------------------- # # Top-level blueprints # # --------------------------------------------------------------------------- # @dataclass(frozen=True) class SingleArmBlueprint: """Author a single-arm dataset in ~15 lines. Given the SingleArm spec, the blueprint derives: state_keys = (joint_state_col,) or (joint_state_col, gripper_state_col) state_dims = (dof,) or (dof, 1) [merged when cols are identical] action_keys = (joint_action_col,) or (joint_action_col, gripper_action_col) action_dims = (dof,) or (dof, 1) delta_mask = arm_mode × dof + gripper_mode × 1 gripper_action_dims = (dof,) arm_layout = ArmLayoutSpec(SINGLE, arm_dof=dof, gripper_index_in_raw=dof) """ schema_id: str robot_type: str arm: SingleArm cameras: Mapping[str, str] annotations: tuple[AnnotationSpec, ...] = () source_path: Optional[str] = None # Consumed POST-build by `discover_schema` (only relaxes the reverse # "info.json has extra cameras" check). Blueprint authoring does NOT read # this flag — it only forwards the value into the built DatasetSchema. allow_extra_cameras: bool = False # Physical semantics of the gripper action channel — see # DatasetSchema.gripper_semantic. Forwarded into the built schema; default # None means no cross-dataset semantic guard. gripper_semantic: Optional[str] = None # Virtual state columns — see DatasetSchema.virtual_state_sources. # Forwarded verbatim; validate_schema enforces the "virtual." prefix and # membership rules on the built schema. virtual_state_sources: Mapping[str, str] = field(default_factory=dict) def build(self) -> DatasetSchema: a = self.arm # Same column for joint + gripper (e.g. oxe_auge observation.joints # holds [j0..j6, grip] in a single column) → collapse to one key. state_collapsed = a.joint_state_col == a.gripper_state_col action_collapsed = a.joint_action_col == a.gripper_action_col if state_collapsed: state_keys = (a.joint_state_col,) state_dims = (a.dof + 1,) else: state_keys = (a.joint_state_col, a.gripper_state_col) state_dims = (a.dof, 1) if action_collapsed: action_keys = (a.joint_action_col,) action_dims = (a.dof + 1,) else: action_keys = (a.joint_action_col, a.gripper_action_col) action_dims = (a.dof, 1) delta_mask = ( tuple([a.arm_mode == "delta"] * a.dof) + (a.gripper_mode == "delta",) ) gripper_action_dims = (a.dof,) arm_layout = ArmLayoutSpec( arm_count=ArmCount.SINGLE, arm_dof=a.dof, gripper_index_in_raw=a.dof, ) annotation_losses = tuple(s.to_loss_spec() for s in self.annotations) image_mapping = expand_camera_mapping(self.cameras) return DatasetSchema( schema_id=self.schema_id, robot_type=self.robot_type, state_keys=state_keys, action_keys=action_keys, state_dims=state_dims, action_dims=action_dims, delta_mask=delta_mask, gripper_action_dims=gripper_action_dims, image_mapping=image_mapping, source="manifest", source_path=self.source_path, allow_extra_cameras=self.allow_extra_cameras, arm_layout=arm_layout, annotation_losses=annotation_losses, gripper_semantic=self.gripper_semantic, virtual_state_sources=self.virtual_state_sources, ) @dataclass(frozen=True) class MultiChannelBlueprint: """Author a dataset whose state/action spans multiple independent channels. Use when the robot-native layout isn't ``arm + 1 gripper`` — e.g. robocoin merges 16+ source robots with a dedicated dual-arm `gripper_open_scale_*` channel separate from the main `action` column. Because `channels` may not map onto any standard ArmLayout, the caller passes `arm_layout_canonical` explicitly (or leaves it None if the dataset is not "arm-shaped" at all). """ schema_id: str robot_type: str channels: tuple[StateActionChannel, ...] cameras: Mapping[str, str] annotations: tuple[AnnotationSpec, ...] = () arm_layout_canonical: Optional[ArmLayoutSpec] = None source_path: Optional[str] = None # Consumed POST-build by `discover_schema`. See SingleArmBlueprint # for the full note; same semantics apply here. allow_extra_cameras: bool = False # See SingleArmBlueprint.gripper_semantic. gripper_semantic: Optional[str] = None def __post_init__(self) -> None: if not self.channels: raise ValueError("MultiChannelBlueprint.channels must be non-empty") def build(self) -> DatasetSchema: state_keys = tuple(c.state_col for c in self.channels) action_keys = tuple(c.action_col for c in self.channels) state_dims = tuple(c.dim for c in self.channels) action_dims = tuple(c.dim for c in self.channels) # delta_mask: concatenate per-channel (dim × mode flag). delta_mask: tuple[bool, ...] = () gripper_idx_list: list[int] = [] offset = 0 for c in self.channels: delta_mask = delta_mask + tuple([c.mode == "delta"] * c.dim) if c.is_gripper: gripper_idx_list.extend(range(offset, offset + c.dim)) offset += c.dim gripper_action_dims = tuple(gripper_idx_list) annotation_losses = tuple(s.to_loss_spec() for s in self.annotations) image_mapping = expand_camera_mapping(self.cameras) return DatasetSchema( schema_id=self.schema_id, robot_type=self.robot_type, state_keys=state_keys, action_keys=action_keys, state_dims=state_dims, action_dims=action_dims, delta_mask=delta_mask, gripper_action_dims=gripper_action_dims, image_mapping=image_mapping, source="manifest", source_path=self.source_path, allow_extra_cameras=self.allow_extra_cameras, arm_layout=self.arm_layout_canonical, annotation_losses=annotation_losses, gripper_semantic=self.gripper_semantic, )