| 'Public ActionCodec2 processor and constrained-generation grammars.' |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from collections.abc import Sequence |
| from numbers import Integral |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| from transformers.processing_utils import ProcessorMixin |
|
|
| from ._validation import active_dimensions |
| from .bpe import BPE_NAME, BPESerializer, SetBPESerializer, serializer_class |
| from .cache import BoundedCache |
| from .configuration import CONFIG_NAME, MODEL_TYPE, ActionCodec2Config |
| from .grammar import ActionCodec2BPEGrammar as ActionCodec2BPEGrammar |
| from .grammar import ActionCodec2Grammar as ActionCodec2Grammar |
| from .hf import ( |
| resolve_pretrained_directory, |
| save_huggingface_metadata, |
| upload_pretrained_directory, |
| ) |
| from .physical_codec import PhysicalCodec |
| from .primitives import PRIMITIVES_NAME, PrimitiveStore |
|
|
|
|
| class ActionCodec2Processor(ProcessorMixin): |
| 'Hugging Face-style processor for continuous robot action trajectories.' |
|
|
| model_type = MODEL_TYPE |
| name = "actioncodec2" |
| is_exact = False |
| set_bpe_is_lossless = True |
| model_input_names = ["input_ids"] |
| attributes = [] |
|
|
| def __init__( |
| self, |
| config: ActionCodec2Config, |
| primitives: PrimitiveStore, |
| serializer: SetBPESerializer | BPESerializer, |
| ) -> None: |
| if serializer.scheme != config.set_bpe.scheme: |
| raise ValueError("serializer scheme does not match config.yaml") |
| if tuple(serializer.bins) != tuple(primitives.bins): |
| raise ValueError("BPE bins do not match primitives.npz") |
| if serializer.lookahead != config.set_bpe.lookahead: |
| raise ValueError("BPE lookahead does not match config.yaml") |
|
|
|
|
| self.config = config |
| self.primitives = primitives |
| self.serializer = serializer |
| self._grammar_cache = BoundedCache(8) |
| self.physical = PhysicalCodec(config, primitives) |
| self.action_dim = config.physical.input_dim |
| self.bins = primitives.bins |
| self.vocab_size = serializer.vocab_size |
| self.lookahead = serializer.lookahead |
| self.requires_previous_first_order = self.physical.requires_previous_first_order |
| zero_cells = [] |
| for dimension, axis in enumerate(primitives.values): |
| matches = np.flatnonzero(np.asarray(axis) == 0) |
| if not len(matches): |
| raise ValueError( |
| "ActionCodec2 artifacts require an exact zero primitive in every axis; " |
| f"dimension {dimension} has none" |
| ) |
| zero_cells.append(int(matches[0])) |
| self._zero_cells = np.asarray(zero_cells, dtype=np.int64) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| pretrained_model_name_or_path: str | Path, |
| **kwargs: Any, |
| ) -> ActionCodec2Processor: |
| 'Load a local directory or Hugging Face Hub snapshot.' |
| directory = resolve_pretrained_directory(pretrained_model_name_or_path, **kwargs) |
| config = ActionCodec2Config.from_pretrained(directory) |
| primitives = PrimitiveStore.from_pretrained(directory, input_dim=config.physical.input_dim) |
| serializer = serializer_class(config).from_pretrained( |
| directory, |
| bins=primitives.bins, |
| lookahead=config.set_bpe.lookahead, |
| ) |
| return cls(config, primitives, serializer) |
|
|
| def save_pretrained( |
| self, |
| save_directory: str | Path, |
| push_to_hub: bool = False, |
| **kwargs: Any, |
| ) -> list[str]: |
| """Save numerical artifacts plus standard AutoProcessor metadata.""" |
| if not push_to_hub and kwargs: |
| raise TypeError(f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}") |
| directory = Path(save_directory) |
| self.config.save_pretrained(directory) |
| self.primitives.save_pretrained(directory) |
| self.serializer.save_pretrained(directory) |
| metadata = save_huggingface_metadata(directory) |
| files = [ |
| directory / CONFIG_NAME, |
| directory / PRIMITIVES_NAME, |
| directory / BPE_NAME, |
| *metadata, |
| ] |
| if push_to_hub: |
| upload_pretrained_directory(directory, **kwargs) |
| return [str(path) for path in files] |
|
|
| def to_dict(self, legacy_serialization: bool = True) -> dict[str, object]: |
| """Return lightweight processor metadata, never large primitive/BPE arrays.""" |
| return { |
| "processor_class": self.__class__.__name__, |
| "model_type": self.model_type, |
| "action_dim": self.action_dim, |
| "vocab_size": self.vocab_size, |
| } |
|
|
| @property |
| def artifact_fingerprint(self) -> str: |
| """Content hash over config, primitive axes, and ordered BPE rules.""" |
| digest = hashlib.sha256() |
| digest.update( |
| json.dumps(self.config.to_dict(), sort_keys=True, separators=(",", ":")).encode() |
| ) |
| for axis in self.primitives.values: |
| values = np.ascontiguousarray(axis) |
| digest.update(str(values.dtype).encode()) |
| digest.update(np.asarray(values.shape, dtype=np.int64).tobytes()) |
| digest.update(values.tobytes()) |
| digest.update( |
| json.dumps(self.serializer.to_dict(), sort_keys=True, separators=(",", ":")).encode() |
| ) |
| return digest.hexdigest() |
|
|
| def __call__( |
| self, actions: np.ndarray, *, previous_first_order: np.ndarray | None = None |
| ) -> dict[str, list[list[int]]]: |
| |
| values = np.asarray(actions) |
| if values.ndim == 2: |
| values = values[None] |
| if previous_first_order is not None: |
| previous_first_order = np.asarray(previous_first_order)[None] |
| return {"input_ids": self.encode_batch(values, previous_first_order=previous_first_order)} |
|
|
| def quantize_batch( |
| self, actions: np.ndarray, *, previous_first_order: np.ndarray | None = None |
| ) -> np.ndarray: |
| """Quantize ``[B,T,D]`` using a ``[B,D]`` first-order boundary if required.""" |
| return self.physical.encode_batch(actions, previous_first_order=previous_first_order) |
|
|
| def _dense_previous_first_order( |
| self, previous_first_order: np.ndarray | None, dims: tuple[int, ...], batch: int |
| ) -> np.ndarray | None: |
| |
| active = set(dims) |
| for component in self.config.physical.components.values(): |
| present = active.intersection(component.dims) |
| if ( |
| component.primitive_order == 2 |
| and component.space == "so3" |
| and present |
| and present != set(component.dims) |
| ): |
| raise ValueError( |
| "second-order SO(3) requires all three active component dimensions" |
| ) |
| if previous_first_order is None: |
| if any( |
| c.primitive_order == 2 and active.intersection(c.dims) |
| for c in self.config.physical.components.values() |
| ): |
| raise ValueError( |
| f"second-order tokenization requires previous_first_order with shape {(batch, len(dims))} in canonical input units" |
| ) |
| return ( |
| np.zeros((batch, self.action_dim), dtype=np.float32) |
| if self.requires_previous_first_order |
| else None |
| ) |
| values = np.asarray(previous_first_order, dtype=np.float32) |
| if values.shape != (batch, len(dims)) or not np.all(np.isfinite(values)): |
| raise ValueError(f"previous_first_order must be finite with shape {(batch, len(dims))}") |
| dense = np.zeros((batch, self.action_dim), dtype=np.float32) |
| dense[:, dims] = values |
| return dense |
|
|
| def tokens_from_cells( |
| self, |
| cells: np.ndarray, |
| *, |
| active_dims: Sequence[int] | None = None, |
| ) -> list[int]: |
| """Apply the lossless tokenizer to one compact or dense cell trajectory.""" |
| return self.serializer.encode(cells, active_dims=active_dims) |
|
|
| def configure_encoding( |
| self, |
| *, |
| spatial_search: bool = False, |
| temporal_search: bool = False, |
| search_window: int = 32, |
| search_beam: int = 4, |
| ) -> ActionCodec2Processor: |
| 'Opt into lossless packing search; both stages default to disabled.' |
| self.serializer.configure_encoding( |
| spatial_search=spatial_search, |
| temporal_search=temporal_search, |
| search_window=search_window, |
| search_beam=search_beam, |
| ) |
| return self |
|
|
| def cells_from_tokens( |
| self, |
| tokens: Sequence[int], |
| *, |
| active_dims: Sequence[int] | None = None, |
| ) -> np.ndarray: |
| """Losslessly recover dense cells with ``-1`` in absent dimensions.""" |
| return self.serializer.decode(tokens, active_dims=active_dims) |
|
|
| def _fill_inactive_cells(self, cells: np.ndarray) -> np.ndarray: |
| |
| values = np.asarray(cells, dtype=np.int64) |
| return np.where(values >= 0, values, self._zero_cells) |
|
|
| def encode( |
| self, |
| actions: np.ndarray, |
| *, |
| active_dims: Sequence[int] | None = None, |
| previous_first_order: np.ndarray | None = None, |
| ) -> list[int]: |
| 'Encode ``[T,D_active]`` with an optional compact ``[D_active]`` boundary.' |
| values = np.asarray(actions, dtype=np.float32) |
| if values.ndim != 2 or len(values) < 1 or values.shape[1] < 1: |
| raise ValueError("actions must have non-empty shape [T, D_active]") |
| dims = ( |
| tuple(range(values.shape[1])) |
| if active_dims is None and values.shape[1] < self.action_dim |
| else active_dimensions(active_dims, self.action_dim) |
| ) |
| if len(dims) != values.shape[1]: |
| raise ValueError("active_dims must have one entry per action coordinate") |
| previous = None if previous_first_order is None else np.asarray(previous_first_order)[None] |
| return self.encode_batch(values[None], active_dims=dims, previous_first_order=previous)[0] |
|
|
| def encode_batch( |
| self, |
| actions: np.ndarray, |
| *, |
| active_dims: Sequence[int] | None = None, |
| previous_first_order: np.ndarray | None = None, |
| ) -> list[list[int]]: |
| 'Encode a dense or compact ``[B,T,D]`` batch into ragged streams.' |
| values = np.asarray(actions, dtype=np.float32) |
| if values.ndim != 3 or len(values) < 1 or values.shape[1] < 1 or values.shape[2] < 1: |
| raise ValueError("actions must have non-empty shape [B, T, D_active]") |
| dims = active_dimensions(active_dims, self.action_dim) |
| if len(dims) != values.shape[2]: |
| raise ValueError("active_dims must have one entry per action coordinate") |
| full_width = dims == tuple(range(self.action_dim)) |
| if full_width: |
| dense = values |
| else: |
| dense = np.zeros((len(values), values.shape[1], self.action_dim), dtype=np.float32) |
| dense[:, :, dims] = values |
| previous = self._dense_previous_first_order(previous_first_order, dims, len(values)) |
| cells = self.quantize_batch(dense, previous_first_order=previous) |
| if not full_width: |
| cells = np.take(cells, dims, axis=2) |
| return self.serializer.encode_batch(cells, active_dims=dims) |
|
|
| def decode( |
| self, |
| tokens: Sequence[int], |
| *, |
| active_dims: Sequence[int] | None = None, |
| previous_first_order: np.ndarray | None = None, |
| ) -> np.ndarray: |
| 'Decode one stream to ``[T,D_active]`` reconstructed actions.' |
| previous = None if previous_first_order is None else np.asarray(previous_first_order)[None] |
| return self.decode_batch([tokens], active_dims=active_dims, previous_first_order=previous)[ |
| 0 |
| ] |
|
|
| def decode_batch( |
| self, |
| streams: Sequence[Sequence[int]], |
| *, |
| active_dims: Sequence[int] | None = None, |
| previous_first_order: np.ndarray | None = None, |
| ) -> np.ndarray: |
| 'Decode equal-length streams into one batched physical array.' |
| dims = active_dimensions(active_dims, self.action_dim) |
| stacked_cells = self.serializer.decode_batch(streams, active_dims=dims) |
| dense_cells = self._fill_inactive_cells(stacked_cells) |
| previous = self._dense_previous_first_order(previous_first_order, dims, len(stacked_cells)) |
| decoded = self.physical.decode_batch(dense_cells, previous_first_order=previous) |
| return np.take(decoded, dims, axis=2) |
|
|
| def final_first_order( |
| self, |
| streams: Sequence[Sequence[int]], |
| *, |
| previous_first_order: np.ndarray, |
| active_dims: Sequence[int] | None = None, |
| executed_steps: int | None = None, |
| ) -> np.ndarray: |
| 'Return the next compact ``[B,D_active]`` reconstructed boundary.' |
| dims = active_dimensions(active_dims, self.action_dim) |
| cells = [self.cells_from_tokens(stream, active_dims=dims) for stream in streams] |
| if not cells or len({len(row) for row in cells}) != 1: |
| raise ValueError("streams must be non-empty and have equal decoded horizons") |
| if executed_steps is not None: |
| if ( |
| isinstance(executed_steps, bool) |
| or not isinstance(executed_steps, Integral) |
| or not 1 <= executed_steps <= len(cells[0]) |
| ): |
| raise ValueError("executed_steps must be within the decoded horizon") |
| cells = [row[:executed_steps] for row in cells] |
| previous = self._dense_previous_first_order(previous_first_order, dims, len(cells)) |
| result = self.physical.final_first_order( |
| self._fill_inactive_cells(np.stack(cells)), previous_first_order=previous |
| ) |
| return np.take(result, dims, axis=1) |
|
|
| def eval(self) -> ActionCodec2Processor: |
| """Return ``self``; the rule processor has no train/eval state.""" |
| return self |
|
|
| def to(self, device: object) -> ActionCodec2Processor: |
| """Return ``self``; all runtime arrays are CPU NumPy values.""" |
| del device |
| return self |
|
|
| def grammar( |
| self, |
| horizon: int, |
| *, |
| active_dims: Sequence[int] | None = None, |
| ) -> ActionCodec2Grammar | ActionCodec2BPEGrammar: |
| """Build a constrained grammar for one horizon and route support.""" |
| if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1: |
| raise ValueError("horizon must be a positive integer") |
| dims = active_dimensions(active_dims, self.action_dim) |
| key = (horizon, dims) |
| grammar = self._grammar_cache.get(key) |
| if grammar is None: |
| if self.serializer.scheme == "unified_set_bpe": |
| from .unified_set_bpe import UnifiedSetBPEGrammar |
|
|
| grammar_class = UnifiedSetBPEGrammar |
| else: |
| grammar_class = ( |
| ActionCodec2BPEGrammar |
| if isinstance(self.serializer, BPESerializer) |
| else ActionCodec2Grammar |
| ) |
| grammar = grammar_class(self.serializer, horizon=horizon, active_dims=dims) |
| self._grammar_cache[key] = grammar |
| return grammar |
|
|
