"""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 ...backends.actioncodec2 import ( ActionCodec2Router, PresetConfig, ) from ...backends.actioncodec2.native_backend import native_kernel from ...backends.actioncodec2.semantic import resampled_horizon from .._types import ActionIdxBatch, TokenIds from ..base_action_tokenizer import BaseActionTokenizer from ..layout import ActionSpaceMixin from ..conversion import _numpy, _numpy_token_rows, restore_source_layout class ActionCodec2Tokenizer(ActionSpaceMixin, BaseActionTokenizer): 'ActionCodec2 frontend bound to physical tensor semantics, not a horizon.' 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) if preset.target_hz != 15.0: raise ValueError( f"ActionCodec2 frontend artifacts require target_hz=15; got {preset.target_hz}. " "Pass the source data rate via fps; use a physical profile calibrated at 15 Hz. " "Changing target_hz alone does not recalibrate physical bounds." ) 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 = bool(backend.second_order_boundaries(preset)) self._encoded_action_dim = sum( len(component.active_dims) for component in preset.components ) self.max_action_token_len = None self.fixed_action_token_len = None super().__init__(append_bar_sentinels=False) @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.' router = ActionCodec2Router.from_pretrained(pretrained_model_name_or_path, **kwargs) return cls._from_router(router, action_space=action_space) @classmethod def _from_router(cls, router: ActionCodec2Router, *, action_space): 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 configure_encoding( self, *, spatial_search: bool = False, temporal_search: bool = False, search_window: int = 32, search_beam: int = 4, ) -> ActionCodec2Tokenizer: 'Opt into native packing search; both independent stages default off.' self.action_tokenizer.configure_encoding( spatial_search=spatial_search, temporal_search=temporal_search, search_window=search_window, search_beam=search_beam, ) return self 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 ) 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 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 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.' 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.' 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: 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.' 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, }, ) def decode_token_ids_to_actions_batch( self, action_token_ids: Sequence[TokenIds], *, time_horizon: int | None = None, action_dim: int | None = None, decode_kwargs: dict | None = None, ) -> torch.Tensor: 'Decode integer rows to ``(B,T,D)``; the native router validates ranges.' if native_kernel("route_decode_batch") is None: return super().decode_token_ids_to_actions_batch( action_token_ids, time_horizon=time_horizon, action_dim=action_dim, decode_kwargs=decode_kwargs, ) actions, _ = self._decode_action_indices( _numpy_token_rows(action_token_ids), time_horizon=time_horizon, action_dim=action_dim, decode_kwargs=decode_kwargs, ) return torch.from_numpy(actions) def decode_prefix( self, action_token_ids, *, horizon: int, steps: int, fps: float | None = None, current_state=None, previous_first_order=None, ) -> torch.Tensor: 'Decode a covered streaming prefix to ``(B,steps,D)`` on the 15 Hz grid.' if self.action_space.input == "absolute" and current_state is None: raise ValueError("absolute prefix decode requires current_state") if self.action_space.input == "delta" and current_state is not None: raise ValueError("delta prefix decode must omit current_state") decoded = self.action_tokenizer.decode_prefix_batch( [self._checked_decode_row(row).tolist() for row in action_token_ids], preset=self.action_space.name, horizon=horizon, steps=steps, fps=float(self._source_fps(fps)), current_state=_numpy(current_state), previous_first_order=self._numpy_boundary(previous_first_order), ) return torch.from_numpy(self._restore_source_layout(decoded)) 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.' 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: 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]]]: 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: return restore_source_layout(self.action_space, decoded)