| """Hugging Face processor over the shared ActionCodec2 frontend implementation.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from collections.abc import Mapping |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| from transformers.feature_extraction_utils import BatchFeature |
| from transformers.processing_utils import ProcessorMixin |
|
|
| from ...backends.actioncodec2 import ActionCodec2Router, PresetConfig |
| from ...backends.actioncodec2.defaults import ( |
| CodecBuilder, |
| ProfileBundleBuilder, |
| builtin_action_spaces, |
| ) |
| from ...backends.actioncodec2.hf import ( |
| PROCESSOR_CONFIG_NAME, |
| ROUTER_ALLOW_PATTERNS, |
| ROUTER_MANIFEST_PATTERNS, |
| ActionCodec2FrontendHFConfig, |
| resolve_pretrained_directory, |
| save_huggingface_metadata, |
| save_model_card, |
| upload_pretrained_directory, |
| ) |
| from ...backends.actioncodec2.preprocessing import PreprocessingSpec |
| from ...backends.actioncodec2.release import MANIFEST_NAME |
| from ...exporting import save_runtime |
| from .loading import describe_release, load_release_router |
| from .tokenizer import ActionCodec2Tokenizer |
| from ..training import TrainingMixin |
| from ...backends.actioncodec2.training import ProfileTrainer |
|
|
|
|
| def _saved_metadata(directory: Path) -> dict: |
| metadata = directory / PROCESSOR_CONFIG_NAME |
| return json.loads(metadata.read_text(encoding="utf-8")) if metadata.is_file() else {} |
|
|
|
|
| def _fitted_status(profiles) -> str: |
| status = ", ".join(sorted(profiles)) or "none; call codec.fit(episodes)" |
| return f"Fitted profiles: {status}\n" |
|
|
|
|
| class ActionCodec2(TrainingMixin, ActionCodec2Tokenizer, ProcessorMixin): |
| 'Action-space-bound processor with lossless token serialization.' |
|
|
| trainer_class = ProfileTrainer |
| model_type = ActionCodec2FrontendHFConfig.model_type |
| model_input_names = ["input_ids"] |
|
|
| attributes = [] |
|
|
| def __init__( |
| self, |
| backend: ActionCodec2Router | None = None, |
| *, |
| action_space: str, |
| primitive_order: int = 1, |
| bpe_scheme: str | None = None, |
| profile_configs: Mapping[str, object] | None = None, |
| ) -> None: |
| fresh = backend is None |
| if fresh: |
| backend = CodecBuilder( |
| primitive_order=primitive_order, |
| bpe_scheme=bpe_scheme, |
| profile_configs=profile_configs, |
| ).build() |
| elif primitive_order != 1 or bpe_scheme is not None or profile_configs is not None: |
| raise ValueError("profile construction options cannot accompany an existing backend") |
| super().__init__(backend, action_space=action_space) |
| ProcessorMixin.__init__(self) |
| self._fitted_profiles = frozenset() if fresh else frozenset(backend.profiles) |
| self.training_report = None |
| self.preprocessing: dict[str, PreprocessingSpec] = {} |
| self.release_metadata: dict[str, Any] | None = None |
|
|
| def __dir__(self): |
| inapplicable = { |
| "audio_tokenizer", |
| "chat_template", |
| "feature_extractor_class", |
| "apply_chat_template", |
| "post_process_image_text_to_text", |
| "pad_action_id", |
| "bos_blk_id", |
| "parts_meta", |
| "BAR_SENTINEL_TOKENS", |
| } |
| return [name for name in super().__dir__() if name not in inapplicable] |
|
|
| def __repr__(self) -> str: |
| return ( |
| f"ActionCodec2(action_space={self.action_space.name!r}, " |
| f"D={self.action_dim}, vocab={self.action_tokenizer.vocab_size})" |
| ) |
|
|
| @staticmethod |
| def builtin_action_spaces(): |
| """Return the immutable named layout definitions shipped in the package.""" |
| return builtin_action_spaces() |
|
|
| @classmethod |
| def from_profiles(cls, joint, eef, *, action_space: str, token_budgets=None): |
| 'Bind an artifact assembled from existing joint/EEF profile directories.' |
| router = ProfileBundleBuilder(joint, eef, token_budgets=token_budgets).build() |
| return cls(router, action_space=action_space) |
|
|
| def logits_processor( |
| self, |
| horizon: int, |
| *, |
| fps: float, |
| prompt_length: int, |
| eos_token_id: int, |
| token_offset: int = 0, |
| pad_token_id=None, |
| ): |
| 'Create a Transformers LogitsProcessor for a fixed action horizon.' |
| from .generation import ActionCodec2LogitsProcessor |
|
|
| return ActionCodec2LogitsProcessor( |
| self, |
| horizon, |
| fps=fps, |
| prompt_length=prompt_length, |
| eos_token_id=eos_token_id, |
| token_offset=token_offset, |
| pad_token_id=pad_token_id, |
| ) |
|
|
| def batch_decode( |
| self, token_rows, *, fps=None, current_state=None, previous_first_order=None, errors="raise" |
| ): |
| 'Decode independent token rows, including different valid horizons.' |
| self._require_fitted() |
| if errors not in ("raise", "return"): |
| raise ValueError("errors must be 'raise' or 'return'") |
| rows = list(token_rows) |
| current = ( |
| None |
| if current_state is None |
| else np.asarray( |
| current_state.detach().cpu() |
| if isinstance(current_state, torch.Tensor) |
| else current_state |
| ) |
| ) |
| if current is not None and current.shape != (len(rows), self.action_dim): |
| raise ValueError(f"current_state must have shape {(len(rows), self.action_dim)}") |
| previous = self._numpy_boundary(previous_first_order) |
| if previous is not None and any( |
| np.ndim(v) != 2 or len(v) != len(rows) for v in previous.values() |
| ): |
| raise ValueError( |
| "previous_first_order arrays must have batch dimension equal to token rows" |
| ) |
| result = [] |
| for index, row in enumerate(rows): |
| try: |
| decoded = self.decode( |
| [row], |
| fps=fps, |
| current_state=None if current is None else current[index : index + 1], |
| previous_first_order=None |
| if previous is None |
| else {k: v[index : index + 1] for k, v in previous.items()}, |
| ) |
| except ValueError as error: |
| error = ValueError(f"token row {index}: {error}") |
| if errors == "raise": |
| raise error from None |
| decoded = error |
| result.append(decoded) |
| return result |
|
|
| def apply_chat_template(self, *args, **kwargs): |
| raise NotImplementedError( |
| "ActionCodec2 accepts action tensors; use codec.encode(actions, fps=...) or codec(actions, fps=...)" |
| ) |
|
|
| def post_process_image_text_to_text(self, *args, **kwargs): |
| raise NotImplementedError( |
| "ActionCodec2 returns actions; use codec.decode(token_rows, fps=...)" |
| ) |
|
|
| @property |
| def fitted_profiles(self) -> frozenset[str]: |
| """Names of profiles fitted on data; persisted by save_pretrained().""" |
| return self._fitted_profiles |
|
|
| @property |
| def is_fitted(self) -> bool: |
| """Whether every profile used by the selected action space is fitted.""" |
| return all(c.space in self.fitted_profiles for c in self.action_space.components) |
|
|
| def _require_fitted(self): |
| missing = {c.space for c in self.action_space.components} - self.fitted_profiles |
| if missing: |
| raise RuntimeError( |
| f"Unfitted profiles {sorted(missing)}; call codec.fit(episodes) first" |
| ) |
|
|
| def _describe_action_spaces(self, action_dim: int | None = None) -> str: |
| return _fitted_status(self.fitted_profiles) + super()._describe_action_spaces(action_dim) |
|
|
| def for_action_space(self, name: str) -> ActionCodec2: |
| """Share current weights with a new binding; subsequent fit replaces only that instance.""" |
| codec = super().for_action_space(name) |
| codec._fitted_profiles = self.fitted_profiles |
| codec.training_report = self.training_report |
| codec.preprocessing = dict(self.preprocessing) |
| codec.release_metadata = self.release_metadata |
| return codec |
|
|
|
|
| def _encode_action_indices(self, action, encode_kwargs=None): |
| self._require_fitted() |
| return super()._encode_action_indices(action, encode_kwargs) |
|
|
| def _decode_action_indices(self, action_indices, **kwargs): |
| self._require_fitted() |
| return super()._decode_action_indices(action_indices, **kwargs) |
|
|
| def decode_prefix(self, action_token_ids, **kwargs): |
| self._require_fitted() |
| return super().decode_prefix(action_token_ids, **kwargs) |
|
|
| def grammar(self, horizon: int, *, fps: float | None = None): |
| self._require_fitted() |
| return super().grammar(horizon, fps=fps) |
|
|
| def final_first_order(self, action_token_ids, *, previous_first_order, executed_steps=None): |
| self._require_fitted() |
| return super().final_first_order( |
| action_token_ids, |
| previous_first_order=previous_first_order, |
| executed_steps=executed_steps, |
| ) |
|
|
| @classmethod |
| def describe_pretrained( |
| cls, |
| pretrained_model_name_or_path: str | Path, |
| *, |
| action_dim: int | None = None, |
| gripper_mode: str = "2bin", |
| primitive_order: int = 1, |
| **kwargs: Any, |
| ) -> str: |
| 'Inspect a local/Hub artifact before choosing an action space.' |
| directory = resolve_pretrained_directory( |
| pretrained_model_name_or_path, allow_patterns=ROUTER_MANIFEST_PATTERNS, **kwargs |
| ) |
| if (directory / MANIFEST_NAME).is_file(): |
| return describe_release( |
| directory, |
| action_dim=action_dim, |
| gripper_mode=gripper_mode, |
| primitive_order=primitive_order, |
| ) |
| specs, _ = ActionCodec2Router.read_manifest(directory) |
| metadata = _saved_metadata(directory) |
| fitted = metadata.get("fitted_profiles", [spec.name for spec in specs]) |
| return _fitted_status(fitted) + ActionCodec2Router.describe_pretrained( |
| directory, selected=metadata.get("action_space"), action_dim=action_dim |
| ) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| pretrained_model_name_or_path: str | Path, |
| *, |
| action_space: str | Path | PresetConfig | Mapping[str, object] | None = None, |
| gripper_mode: str | None = None, |
| primitive_order: int | None = None, |
| verify_gates: bool = False, |
| **kwargs: Any, |
| ) -> ActionCodec2: |
| 'Load a local/Hub artifact and restore or override its action space.' |
| if cls is ActionCodec2 and ("variant" in kwargs or (Path(pretrained_model_name_or_path) / "variants.json").is_file() |
| or (Path(pretrained_model_name_or_path) / "artifact_manifest.json").is_file()): |
| from ...catalog import ActionCodec2 as Catalog |
| return Catalog.from_pretrained(pretrained_model_name_or_path, action_space=action_space, **kwargs) |
| directory = resolve_pretrained_directory( |
| pretrained_model_name_or_path, allow_patterns=ROUTER_ALLOW_PATTERNS, **kwargs |
| ) |
| if action_space is None: |
| action_space = _saved_metadata(directory).get("action_space") |
| if action_space is None: |
| raise ValueError( |
| "action_space is required when loading an unbound router artifact. " |
| "First inspect its layouts with print(ActionCodec2.describe_pretrained(path)), " |
| "then pass action_space='NAME'; shape alone cannot identify physical meaning." |
| ) |
| if (directory / MANIFEST_NAME).is_file(): |
| router, preprocessing, metadata = load_release_router( |
| directory, |
| gripper_mode="2bin" if gripper_mode is None else gripper_mode, |
| primitive_order=1 if primitive_order is None else primitive_order, |
| verify_gates=verify_gates, |
| ) |
| codec = cls._from_router(router, action_space=action_space) |
| codec.preprocessing, codec.release_metadata = preprocessing, metadata |
| return codec |
| if gripper_mode is not None or primitive_order is not None or verify_gates: |
| raise ValueError( |
| "gripper_mode, primitive_order and verify_gates select release variants only" |
| ) |
| codec = super().from_pretrained(directory, action_space=action_space) |
| metadata = _saved_metadata(directory) |
| fitted = metadata.get("fitted_profiles", list(codec.action_tokenizer.profiles)) |
| if not isinstance(fitted, list) or any( |
| not isinstance(name, str) or name not in codec.action_tokenizer.profiles |
| for name in fitted |
| ): |
| raise ValueError("invalid fitted_profiles in processor_config.json") |
| codec._fitted_profiles = frozenset(fitted) |
| policies = metadata.get("preprocessing", {}) |
| if ( |
| not isinstance(policies, dict) |
| or policies.keys() - codec.action_tokenizer.profiles.keys() |
| ): |
| raise ValueError("invalid preprocessing profiles in processor_config.json") |
| codec.preprocessing = { |
| name: PreprocessingSpec.from_dict(spec) for name, spec in policies.items() |
| } |
| codec.release_metadata = metadata.get("release") |
| report_path = directory / "fit_report.json" |
| if report_path.is_file(): |
| codec.training_report = json.loads(report_path.read_text(encoding="utf-8")) |
| return codec |
|
|
| def to_dict(self, legacy_serialization: bool = True) -> dict[str, object]: |
| 'Return lightweight discovery metadata.' |
| del legacy_serialization |
| payload = { |
| "processor_class": type(self).__name__, |
| "model_type": self.model_type, |
| "action_space": self.action_space.name, |
| "fitted_profiles": sorted(self.fitted_profiles), |
| } |
| if self.preprocessing: |
| payload["preprocessing"] = { |
| name: spec.to_dict() for name, spec in self.preprocessing.items() |
| } |
| return payload |
|
|
| def save_pretrained( |
| self, save_directory: str | Path, push_to_hub: bool = False, *, include_training_report=False, **kwargs: Any |
| ) -> list[str]: |
| """Save the complete router and the selected action-space binding.""" |
| if not push_to_hub and kwargs: |
| raise TypeError(f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}") |
| directory = Path(save_directory) |
| files = self.action_tokenizer.save_pretrained(directory) |
|
|
| save_model_card(directory, str(self)) |
| save_huggingface_metadata( |
| directory, |
| **self.to_dict(), |
| ) |
| report_path = directory / "fit_report.json" |
| if include_training_report and self.training_report is not None: |
| report_path.write_text( |
| json.dumps(self.training_report, indent=2) + "\n", encoding="utf-8" |
| ) |
| files.append(str(report_path)) |
| elif report_path.exists(): |
| report_path.unlink() |
| files.extend(save_runtime(directory)) |
| if push_to_hub: |
| upload_pretrained_directory(directory, **kwargs) |
| return files |
|
|
| def __call__( |
| self, |
| action: torch.Tensor | np.ndarray, |
| *, |
| fps: float | None = None, |
| current_state: torch.Tensor | np.ndarray | None = None, |
| previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None, |
| encode_kwargs: dict | None = None, |
| return_tensors: str | None = None, |
| ) -> BatchFeature: |
| |
| rows = self.encode( |
| action, |
| fps=fps, |
| current_state=current_state, |
| previous_first_order=previous_first_order, |
| encode_kwargs=encode_kwargs, |
| ) |
| return BatchFeature({"input_ids": rows}, tensor_type=return_tensors) |
|
|
