"""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 .routing.defaults import CodecBuilder, ProfileBundleBuilder, builtin_action_spaces from .integration.hf import ( PROCESSOR_CONFIG_NAME, ROUTER_ALLOW_PATTERNS, ActionCodec2FrontendHFConfig, resolve_pretrained_directory, save_huggingface_metadata, save_model_card, upload_pretrained_directory, ) from .routing.presets import PresetConfig from .routing.routing import ActionCodec2Router from .frontend.tokenizer import ActionCodec2Tokenizer def _saved_metadata(directory: Path) -> dict[str, Any]: 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(ActionCodec2Tokenizer, ProcessorMixin): """Action-space-bound processor with lossless token serialization. ``encode`` returns ragged token-ID rows; ``decode`` reconstructs a CPU ``torch.float32`` tensor in the original ``(B,T,D)`` column layout. Calling the processor wraps those same rows in a ``BatchFeature`` under ``input_ids``. Physical quantization remains lossy as configured. Released first- and second-order bundles both register 13 layouts: single and dual EEF arms, single and dual arms with 6 or 7 joint coordinates, each with absolute/delta variants and binary grippers, plus a six-continuous- coordinate absolute layout without a separate binary gripper column. The registered names and column meanings come from the loaded artifact. ``print(codec)`` lists all of them; ``print(cls.describe_pretrained(path))`` inspects them before binding. ``print_action_spaces(action_dim=D)`` only filters candidates, because shape alone cannot identify physical meaning. Inputs accept ``(T,D)`` or ``(B,T,D)`` with any positive T; unequal-length episodes must be encoded separately. Absolute layouts require the same ``current_state(B,D)`` for encode/decode. Second-order physical profiles additionally require ``previous_first_order: {component: (B,D_active)}`` at the action space's configured codec rate, independently of the absolute/delta choice. Boundary arrays retain B=1 for a single episode. Physical order does not change the caller's action layout into an acceleration tensor. Args: backend: Optional fitted router. Omit to construct a fresh codec for fit(). action_space: Registered name defining the input tensor's semantics. primitive_order: Physical primitive order, 1 or 2, for fresh construction. bpe_scheme: Optional 'set_bpe' or 'bpe' override for fresh construction. profile_configs: Optional joint/EEF physical configs, mappings or YAML paths. Defaults ship inside the package; no repository checkout is required. """ model_type = ActionCodec2FrontendHFConfig.model_type model_input_names = ["input_ids"] # Numerical artifacts are owned by the router, not modality sub-processors. 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._source_directory = Path(__file__).parent def __dir__(self): inapplicable = { "action_tokenizer", "attributes", "audio_tokenizer", "chat_template", "check_argument_for_proper_class", "encode_kwargs", "feature_extractor_class", "from_args_and_dict", "get_possibly_dynamic_module", "get_processor_dict", "model_input_names", "model_type", "optional_attributes", "optional_call_args", "apply_chat_template", "post_process_image_text_to_text", "pad_action_id", "bos_blk_id", "parts_meta", "BAR_SENTINEL_TOKENS", "push_to_hub", "register_for_auto_class", "tokenizer_class", "to_dict", "to_json_file", "to_json_string", "valid_processor_kwargs", "validate_init_kwargs", } 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. Args: joint: Fitted joint profile directory. eef: Fitted EEF profile directory. action_space: Explicit built-in tensor layout to bind. token_budgets: Optional joint/EEF overrides; defaults to recorded training budgets, or fitted vocabulary sizes if unavailable. """ 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. ``prompt_length`` excludes leading model prompt IDs from the grammar; ``token_offset`` explicitly maps codec IDs to a contiguous model range. EOS/padding IDs must be outside that range. No model vocabulary is changed. """ from .tokenization.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. Args: token_rows: Iterable of token-ID rows. Each must cover a complete trajectory; intended generation horizon must be checked separately with grammar.is_complete before decode. fps: Output sampling rate, shared by all rows. current_state: Optional (B,D) preceding states in original layout. previous_first_order: Optional {component: (B,D_active)} boundaries. errors: 'raise' stops at the first invalid row; 'return' returns the ValueError at that position so other valid rows remain usable. Returns: List of CPU float32 tensors shaped (1,T_i,D), or ValueError entries when errors='return'. No truncated row is padded or executed partially. """ 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 @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._source_directory = self._source_directory return codec def fit( self, episodes, *, fps: float | None = None, current_state=None, previous_first_order=None, bootstrap_context: bool = False, backend: str = "auto", threads: int = 0, batch_size: int = 256, progress: bool = False, max_cell_bytes: int | None = None, work_dir: str | Path | None = None, ) -> ActionCodec2: """Fit physical vocabularies from actions in declared source layouts. Args: episodes: Array/tensor (T,D), (B,T,D), or iterable of variable-length episodes. Records may contain ``actions``, ``action_space`` (a registered name), ``fps``, ``current_state(B,D)``, and ``previous_first_order: {component: (B,D_active)}``. A record's metadata must not duplicate the corresponding keyword argument. fps: Source frequency override shared by all input episodes. current_state: Preceding source state (B,D), required for absolute inputs. Use per-record values when episodes have different states. previous_first_order: Preceding canonical codec-rate increments by component, with shape (B,D_active), required for second-order profiles. bootstrap_context: Use two measured canonical frames as second-order context, training on the remaining suffix. Mutually exclusive with previous_first_order. Inference still requires explicit boundaries. backend: Exact BPE implementation: 'auto', 'native', or 'python'. threads: Native worker count; 0 uses the backend default. batch_size: Maximum canonical trajectories per quantization batch. progress: Display BPE training progress. max_cell_bytes: Hard bound on retained quantized corpus bytes per profile. work_dir: Parent scratch directory for sequential source staging files, automatically removed on success or failure. Defaults to system temp. Returns: This instance, with per-profile results in ``training_report``. Only profiles present in the data are fitted; other weights are preserved. Vocabulary replacement is atomic on success. Refit changes token IDs: regenerate previously encoded datasets before using the new vocabulary. """ # Training is intentionally lazy: a loaded Hub artifact only needs the # inference runtime. The full fitting pipeline remains available from # the source package and is imported when the user calls fit(). from .training.fitting import ActionCodec2Fitter router, report = ActionCodec2Fitter( self, bootstrap_context=bootstrap_context ).fit( episodes, work_dir=work_dir, metadata=dict( fps=fps, current_state=current_state, previous_first_order=previous_first_order, ), backend=backend, threads=threads, batch_size=batch_size, progress=progress, max_cell_bytes=max_cell_bytes, ) fitted = self.fitted_profiles | report["profiles"].keys() replacement = type(self)(router, action_space=self.action_space.name) replacement._fitted_profiles = frozenset(fitted) replacement.training_report = report replacement._source_directory = self._source_directory replacement._warned_default_fps = self._warned_default_fps self.__dict__.update(replacement.__dict__) return self 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 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, **kwargs: Any, ) -> str: """Inspect a local/Hub artifact before choosing an action space. Args: pretrained_model_name_or_path: Saved processor/router directory or Hugging Face Hub repository ID. Loads its fitted router for inspection. action_dim: Optional D from ``actions.shape[-1]``. Only filters the report; tensor shape cannot establish physical semantics. **kwargs: Hub options such as ``revision``, ``subfolder``, ``cache_dir``, ``token``, and ``local_files_only``. Returns: Human-readable layout guide, suitable for ``print(...)``. No binding is required or inferred, and no action data is encoded. """ directory = resolve_pretrained_directory( pretrained_model_name_or_path, allow_patterns=ROUTER_ALLOW_PATTERNS, **kwargs, ) router = ActionCodec2Router.from_pretrained(directory) metadata = _saved_metadata(directory) return _fitted_status( metadata.get("fitted_profiles", router.profiles) ) + router.describe_action_spaces( 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, **kwargs: Any, ) -> ActionCodec2: """Load a local/Hub artifact and restore or override its action space. Args: pretrained_model_name_or_path: Saved processor or router directory, or a Hugging Face Hub repository ID. action_space: Registered name, preset YAML path, config, or mapping. Required for an unbound router; otherwise defaults to the saved binding. Custom definitions are saved in the router manifest. **kwargs: Hub options including ``revision``, ``subfolder``, ``cache_dir``, ``token``, and ``local_files_only``. """ 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." ) 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) codec._source_directory = directory 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. ``legacy_serialization`` is accepted for ProcessorMixin signature compatibility. ActionCodec2 has one serialization format for both values. """ del legacy_serialization return { "processor_class": type(self).__name__, "model_type": self.model_type, "action_space": self.action_space.name, "fitted_profiles": sorted(self.fitted_profiles), } def save_pretrained( self, save_directory: str | Path, push_to_hub: bool = 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))}" ) from .integration.exporting import ArtifactExporter directory = Path(save_directory) save_model_card(directory, str(self)) files = self.action_tokenizer.save_pretrained(directory, include_code=False) files.extend( ArtifactExporter( directory, source_directory=self._source_directory ).export() ) save_huggingface_metadata( directory, **self.to_dict(), auto_map={"AutoProcessor": "processing_actioncodec2.ActionCodec2"}, ) report_path = directory / "fit_report.json" if 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() 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: """Prepare ``input_ids`` from ``(B,T,D)`` or a single ``(T,D)`` action. ``current_state(B,D)`` is required for absolute actions. Second-order artifacts require ``previous_first_order`` mapping component names to ``(B,D_active)`` canonical increments, as in :meth:`encode`. ``return_tensors`` may be ``"pt"`` or ``"np"`` for equal-length rows; omitted, rows remain ragged lists without adding padding tokens. """ 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)