| """Action-space-bound, arbitrary-horizon ActionCodec2 frontend.""" |
|
|
| from __future__ import annotations |
|
|
| import sys |
| import warnings |
| from collections.abc import Mapping, Sequence |
| from pathlib import Path |
| from typing import TextIO |
|
|
| import numpy as np |
| import torch |
|
|
| from ..common._frontend_base import BaseActionTokenizer |
| from ..common._types import ActionIdxBatch, TokenIds |
| from ..routing.presets import PresetConfig |
| from ..routing.routing import ActionCodec2Router |
| from ..physical.semantic import resampled_horizon |
|
|
|
|
| def _numpy(value): |
| """Convert a tensor to CPU float32 without changing NumPy input semantics.""" |
| return ( |
| value.detach().cpu().float().numpy() |
| if isinstance(value, torch.Tensor) |
| else value |
| ) |
|
|
|
|
| class ActionCodec2Tokenizer(BaseActionTokenizer): |
| """ActionCodec2 frontend bound to physical tensor semantics, not a horizon. |
| |
| The selected action space fixes the meaning of every source column. Sequence |
| length remains unrestricted: encode and decode infer any positive ``T`` |
| directly from the lossless token coverage. |
| |
| Use :meth:`from_pretrained` for the normal one-call loading path. Built-in |
| and registered action spaces are discoverable through :meth:`print_action_spaces` |
| and are included in :func:`str` output. |
| |
| Args: |
| backend: Router containing the shared physical processors. |
| action_space: Registered action-space name to use for this frontend. |
| """ |
|
|
| encode_kwargs = frozenset({"fps", "current_state", "previous_first_order"}) |
|
|
| def __init__(self, backend: ActionCodec2Router, *, action_space: str) -> None: |
| if not isinstance(backend, ActionCodec2Router): |
| raise TypeError( |
| f"ActionCodec2Tokenizer requires ActionCodec2Router, got {type(backend).__name__}" |
| ) |
| preset = backend.preset(action_space) |
| self.action_tokenizer = backend |
| self._action_space = preset |
| self._action_vocab_size = int(backend.vocab_size) |
| self._new_action_tokens = list(backend.new_action_tokens) |
| self.action_dim = preset.input_dim |
| self._warned_default_fps = False |
| self.requires_previous_first_order = any( |
| physical.primitive_order == 2 |
| and set(component.active_dims).intersection(physical.dims) |
| for component in preset.components |
| for physical in backend.profiles[ |
| component.space |
| ].processor.config.physical.components.values() |
| ) |
| self._encoded_action_dim = sum( |
| len(component.active_dims) for component in preset.components |
| ) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| pretrained_model_name_or_path: str | Path, |
| *, |
| action_space: str | Path | PresetConfig | Mapping[str, object], |
| **kwargs: object, |
| ) -> ActionCodec2Tokenizer: |
| """Load a local/Hub router and bind one explicit physical action space. |
| |
| Args: |
| pretrained_model_name_or_path: Local artifact directory or Hub repo ID. |
| action_space: Registered name, preset YAML path, ``PresetConfig``, or |
| strict preset mapping. Custom configs are registered before binding. |
| **kwargs: Standard Hub snapshot options such as ``revision``, |
| ``cache_dir``, ``token``, and ``local_files_only``. |
| """ |
| router = ActionCodec2Router.from_pretrained( |
| pretrained_model_name_or_path, **kwargs |
| ) |
| if isinstance(action_space, (PresetConfig, Mapping, Path)): |
| name = router.register_action_space(action_space).name |
| elif isinstance(action_space, str): |
| candidate = Path(action_space) |
| name = ( |
| router.register_action_space(candidate).name |
| if candidate.is_file() |
| else action_space |
| ) |
| else: |
| raise TypeError( |
| "action_space must be a registered name, PresetConfig, mapping, or YAML path" |
| ) |
| return cls(router, action_space=name) |
|
|
| def __str__(self) -> str: |
| """Show supported tensor layouts, physical units, and invocation requirements.""" |
| return self._describe_action_spaces() |
|
|
| def _describe_action_spaces(self, action_dim: int | None = None) -> str: |
| return self.action_tokenizer.describe_action_spaces( |
| selected=self.action_space.name, action_dim=action_dim |
| ) |
|
|
| @property |
| def action_space(self) -> PresetConfig: |
| """The immutable tensor semantics selected for this tokenizer.""" |
| return self._action_space |
|
|
| @property |
| def action_spaces(self) -> Mapping[str, PresetConfig]: |
| """Every built-in or user-registered option in the loaded artifact.""" |
| return self.action_tokenizer.action_spaces |
|
|
| def print_action_spaces( |
| self, *, action_dim: int | None = None, file: TextIO | None = None |
| ) -> None: |
| """Print all layouts or candidates matching ``action_dim`` without changing the binding.""" |
| description = self._describe_action_spaces(action_dim) |
| print(description, file=sys.stdout if file is None else file) |
|
|
| def register_action_space( |
| self, |
| preset: PresetConfig | Mapping[str, object] | str | Path, |
| ) -> PresetConfig: |
| """Register a custom action space without changing this instance's binding.""" |
| return self.action_tokenizer.register_action_space(preset) |
|
|
| def for_action_space(self, name: str) -> ActionCodec2Tokenizer: |
| """Return a new tokenizer sharing weights but bound to another option.""" |
| return type(self)(self.action_tokenizer, action_space=name) |
|
|
| def _source_fps(self, fps, preset=None): |
| preset = self.action_space if preset is None else preset |
| if fps is None: |
| if not self._warned_default_fps: |
| warnings.warn( |
| f"fps was omitted; assuming {preset.source_hz:g} Hz from action_space " |
| f"{preset.name!r}. Pass the actual fps explicitly to avoid changing action timing.", |
| UserWarning, |
| stacklevel=3, |
| ) |
| self._warned_default_fps = True |
| return preset.source_hz |
| return fps |
|
|
| def first_order_template(self, batch_size: int) -> dict[str, np.ndarray]: |
| """Explicit zero boundary templates, shaped {component: (B,D_active)}. |
| |
| Zeros are valid only when the preceding canonical increment is known to |
| be zero. Otherwise replace them with measured/reconstructed values at the |
| action space's configured codec rate. |
| Only active second-order components are included; first order returns {}. |
| """ |
| if type(batch_size) is not int or batch_size < 1: |
| raise ValueError("batch_size must be a positive integer") |
| return { |
| c.name: np.zeros((batch_size, len(c.active_dims)), dtype=np.float32) |
| for c in self.action_space.components |
| if any( |
| p.primitive_order == 2 and set(c.active_dims).intersection(p.dims) |
| for p in self.action_tokenizer.profiles[ |
| c.space |
| ].processor.config.physical.components.values() |
| ) |
| } |
|
|
| def max_tokens_for_horizon(self, horizon: int, *, fps: float | None = None) -> int: |
| """Return the token bound after converting a source chunk to the codec clock. |
| |
| This helper describes encode input geometry. It does not store |
| ``horizon`` in the token stream or pass it to decode. |
| """ |
| source_fps = self._source_fps(fps) |
| internal = resampled_horizon( |
| horizon, source_hz=source_fps, target_hz=self.action_space.target_hz |
| ) |
| return internal * self._encoded_action_dim |
|
|
| def grammar(self, horizon: int, *, fps: float | None = None): |
| """Build a grammar for the codec steps implied by source chunk geometry. |
| |
| The grammar constrains token coverage only. When ``fps`` exceeds the |
| codec clock, several source row counts can map to the same codec |
| step count; decode therefore returns the unique canonical resampling |
| implied by the tokens and requested ``fps``, not a hidden source count. |
| """ |
| source_fps = self._source_fps(fps) |
| internal = resampled_horizon( |
| horizon, source_hz=source_fps, target_hz=self.action_space.target_hz |
| ) |
| return self.action_tokenizer.grammar( |
| preset=self.action_space.name, horizon=internal |
| ) |
|
|
| 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, |
| ) -> ActionIdxBatch: |
| """Encode arbitrary-length ``(B,T,D)`` or ``(T,D)`` actions. |
| |
| Second order requires ``previous_first_order`` mapping component names |
| to canonical ``(B,D_active)`` increments at the configured codec rate (``B=1`` for a single |
| trajectory). Absolute inputs also require ``current_state(B,D)``. |
| """ |
| kwargs = self._check_encode_kwargs(encode_kwargs) |
| direct = { |
| "fps": fps, |
| "current_state": current_state, |
| "previous_first_order": previous_first_order, |
| } |
| duplicates = sorted( |
| key for key, value in direct.items() if value is not None and key in kwargs |
| ) |
| if duplicates: |
| raise ValueError(f"encode arguments provided twice: {duplicates}") |
| kwargs.update( |
| {key: value for key, value in direct.items() if value is not None} |
| ) |
| reference = kwargs.get("current_state") |
| if self.action_space.input == "absolute" and reference is None: |
| raise ValueError( |
| f"action space {self.action_space.name!r} is absolute; encode requires " |
| "current_state with shape [B, D] to preserve the first transition" |
| ) |
| if self.action_space.input == "delta" and reference is not None: |
| raise ValueError( |
| f"action space {self.action_space.name!r} is delta; omit current_state" |
| ) |
| values = torch.as_tensor(action) |
| if values.ndim == 2: |
| values = values.unsqueeze(0) |
| return self._encode_action_indices(values, encode_kwargs=kwargs) |
|
|
| encode = __call__ |
|
|
| def decode( |
| self, |
| action_token_ids: Sequence[TokenIds], |
| *, |
| fps: float | None = None, |
| current_state: torch.Tensor | np.ndarray | None = None, |
| previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None, |
| ) -> torch.Tensor: |
| """Decode a batch, deriving sequence length from tokens and output rate. |
| |
| At the codec rate, token coverage determines ``T`` exactly. Other output rates |
| deterministically resample that decoded path. No source horizon is |
| accepted because downsampling cannot preserve an arbitrary original row |
| count without encoding extra metadata. |
| ``previous_first_order`` maps component names to canonical |
| ``(B,D_active)`` codec-rate increments; second order requires the same |
| boundary as encoding, independently of ``current_state(B,D)``. |
| """ |
| return self._decode_token_ids_to_actions_batch( |
| action_token_ids, |
| action_dim=self.action_dim, |
| decode_kwargs={ |
| "fps": fps, |
| "current_state": current_state, |
| "previous_first_order": previous_first_order, |
| }, |
| ) |
|
|
| @staticmethod |
| def _numpy_boundary(previous_first_order): |
| if previous_first_order is None: |
| return None |
| if not isinstance(previous_first_order, Mapping): |
| raise ValueError( |
| "previous_first_order must map component names to [B,D_active] arrays" |
| ) |
| return {name: _numpy(value) for name, value in previous_first_order.items()} |
|
|
| def final_first_order( |
| self, |
| action_token_ids: Sequence[TokenIds], |
| *, |
| previous_first_order: Mapping[str, torch.Tensor | np.ndarray], |
| executed_steps: int | None = None, |
| ) -> dict[str, np.ndarray]: |
| """Return reconstructed component states after a canonical-step prefix. |
| |
| Results are ``[B,D_active]`` arrays at the configured codec rate, |
| suitable for the next encode/decode call. ``executed_steps`` counts |
| codec steps. |
| """ |
| return self.action_tokenizer.final_first_order( |
| [self._checked_decode_row(row).tolist() for row in action_token_ids], |
| preset=self.action_space.name, |
| previous_first_order=self._numpy_boundary(previous_first_order), |
| executed_steps=executed_steps, |
| ) |
|
|
| def _encode_action_indices( |
| self, action: torch.Tensor, encode_kwargs: dict | None = None |
| ) -> ActionIdxBatch: |
| """Encode finite ``(B,T,D)`` actions without fixing ``T`` at construction.""" |
| kwargs = dict(encode_kwargs or {}) |
| fps = float(self._source_fps(kwargs.get("fps"))) |
| current_state = _numpy(kwargs.get("current_state")) |
| values = _numpy(action) |
| if values.ndim != 3 or values.shape[2] != self.action_space.input_dim: |
| raise ValueError( |
| f"actions must have shape [B, T_source, {self.action_space.input_dim}]" |
| ) |
| return self.action_tokenizer.encode_batch( |
| values, |
| preset=self.action_space.name, |
| fps=fps, |
| current_state=current_state, |
| previous_first_order=self._numpy_boundary( |
| kwargs.get("previous_first_order") |
| ), |
| ) |
|
|
| def _decode_action_indices( |
| self, |
| action_indices: ActionIdxBatch, |
| *, |
| time_horizon: int | None = None, |
| action_dim: int | None = None, |
| decode_kwargs: dict | None = None, |
| ) -> tuple[np.ndarray, list[set[str]]]: |
| """Decode rows to the selected preset's original ``(B,T,D)`` layout.""" |
| if time_horizon is not None: |
| raise ValueError( |
| "ActionCodec2 decode infers sequence length from tokens; " |
| "do not provide time_horizon" |
| ) |
| if action_dim is not None and int(action_dim) != self.action_dim: |
| raise ValueError( |
| f"action_dim={action_dim} disagrees with action space " |
| f"{self.action_space.name!r} (D={self.action_dim})" |
| ) |
| kwargs = dict(decode_kwargs or {}) |
| unknown = sorted(set(kwargs) - self.encode_kwargs) |
| if unknown: |
| raise ValueError( |
| f"ActionCodec2Tokenizer does not consume decode_kwargs {unknown}" |
| ) |
| raw_fps = kwargs.get("fps") |
| fps = float(self._source_fps(raw_fps)) |
| current_state = _numpy(kwargs.get("current_state")) |
| if self.action_space.input == "absolute" and current_state is None: |
| raise ValueError( |
| f"action space {self.action_space.name!r} is absolute; decode requires " |
| "current_state with shape [B, D] to restore its external reference" |
| ) |
| if self.action_space.input == "delta" and current_state is not None: |
| raise ValueError( |
| f"action space {self.action_space.name!r} is delta; omit current_state " |
| "to decode in the registered input semantics" |
| ) |
| decoded = self.action_tokenizer.decode_batch( |
| action_indices, |
| preset=self.action_space.name, |
| fps=fps, |
| current_state=current_state, |
| previous_first_order=self._numpy_boundary( |
| kwargs.get("previous_first_order") |
| ), |
| ) |
| actions = self._restore_source_layout(decoded) |
| return actions, [set() for _ in range(actions.shape[0])] |
|
|
| def _restore_source_layout(self, decoded: Mapping[str, np.ndarray]) -> np.ndarray: |
| """Place component outputs back into their registered source columns.""" |
| first = next(iter(decoded.values())) |
| batch, steps = first.shape[:2] |
| if any(values.shape[:2] != (batch, steps) for values in decoded.values()): |
| raise RuntimeError( |
| "decoded components disagree on batch or sequence length" |
| ) |
| actions = np.empty((batch, steps, self.action_dim), dtype=np.float32) |
| for component in self.action_space.components: |
| values = decoded[component.name] |
| motion_width = len(component.source_dims) |
| actions[:, :, component.source_dims] = values[:, :, :motion_width] |
| if component.gripper_dim is not None: |
| actions[:, :, component.gripper_dim] = values[:, :, motion_width] |
| return actions |
|
|
