| """Abstract base for frontend action tokenizers.""" |
|
|
| from __future__ import annotations |
|
|
| import logging |
| from abc import ABC, abstractmethod |
| from collections.abc import Sequence |
|
|
| import numpy as np |
| import torch |
|
|
| from ._types import ActionIdxBatch, DecodeResult, TokenIds |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class BaseActionTokenizer(ABC): |
| """Base action tokenizer: maps actions ↔ discrete action-token ids.""" |
|
|
| BAR_SENTINEL_TOKENS = ("<bos_blk>", "<eos_blk>", "<pad_action_token>") |
|
|
| encode_kwargs: frozenset[str] = frozenset() |
|
|
| drops_noop_parts: bool = False |
|
|
| parts_meta: dict[str, int] | None = None |
|
|
| def __init__(self, *, append_bar_sentinels: bool = True) -> None: |
| required = ( |
| "action_tokenizer", |
| "_action_vocab_size", |
| "_new_action_tokens", |
| "max_action_token_len", |
| "fixed_action_token_len", |
| ) |
| missing = [a for a in required if not hasattr(self, a)] |
| if missing: |
| raise ValueError( |
| f"{type(self).__name__}.__init__ must set {missing} before super().__init__()" |
| ) |
| if self.max_action_token_len is not None: |
| self.max_action_token_len = int(self.max_action_token_len) |
| if self.max_action_token_len <= 0: |
| raise ValueError( |
| f"{type(self).__name__}: max_action_token_len must be positive, " |
| f"got {self.max_action_token_len}" |
| ) |
| if self.fixed_action_token_len is not None: |
| self.fixed_action_token_len = int(self.fixed_action_token_len) |
| if self.max_action_token_len is None: |
| raise ValueError( |
| f"{type(self).__name__}: fixed_action_token_len requires a finite " |
| "max_action_token_len" |
| ) |
| if self.fixed_action_token_len != self.max_action_token_len: |
| raise ValueError( |
| f"{type(self).__name__}: fixed_action_token_len=" |
| f"{self.fixed_action_token_len} must equal max_action_token_len=" |
| f"{self.max_action_token_len} — a fixed-length row is its own maximum" |
| ) |
| base = int(self._action_vocab_size) |
| self._new_action_tokens = list(self._new_action_tokens) |
| if append_bar_sentinels: |
| self._new_action_tokens += list(self.BAR_SENTINEL_TOKENS) |
| self.bos_blk_id: int | None = base |
| self.eos_blk_id: int | None = base + 1 |
| self.pad_action_id: int | None = base + 2 |
| self._action_vocab_size = base + len(self.BAR_SENTINEL_TOKENS) |
| else: |
| self.bos_blk_id = self.eos_blk_id = self.pad_action_id = None |
|
|
| @property |
| def new_action_tokens(self) -> list[str]: |
| """Token strings in id order — downstream registers these into its VLM vocab.""" |
| return list(self._new_action_tokens) |
|
|
| def to_token_strs(self, token_ids: TokenIds) -> list[str]: |
| """Render one id row as its token strings, e.g. ``<right_arm_0>`` / ``<action0042>``.""" |
| row = self._to_id_row(token_ids).tolist() |
| n = len(self._new_action_tokens) |
| bad = [i for i in row if not (0 <= i < n)] |
| if bad: |
| raise ValueError(f"token ids {bad[:5]} out of range [0, {n})") |
| return [self._new_action_tokens[i] for i in row] |
|
|
| def to(self, device_id): |
| """Move the backend to ``device_id`` and record it as ``self.device``.""" |
| self.device = torch.device(device_id) |
| self.action_tokenizer = self.action_tokenizer.to(device_id) |
| return self |
|
|
| def _check_encode_kwargs(self, encode_kwargs: dict | None) -> dict: |
| kw = dict(encode_kwargs or {}) |
| unknown = sorted(set(kw) - self.encode_kwargs) |
| if unknown: |
| raise ValueError( |
| f"{type(self).__name__} does not consume encode_kwargs {unknown}; " |
| f"it declares {sorted(self.encode_kwargs)}" |
| ) |
| return kw |
|
|
| def __call__( |
| self, |
| action: torch.Tensor, |
| *, |
| encode_kwargs: dict | None = None, |
| ) -> ActionIdxBatch: |
| return self._encode_action_indices( |
| action, encode_kwargs=self._check_encode_kwargs(encode_kwargs) |
| ) |
|
|
| def _strip_bar_sentinels(self, row: torch.Tensor) -> torch.Tensor: |
| if self.bos_blk_id is None: |
| return row |
| keep = (row != self.bos_blk_id) & (row != self.eos_blk_id) & (row != self.pad_action_id) |
| return row[keep] |
|
|
| @staticmethod |
| def _to_id_row(token_ids: TokenIds) -> torch.Tensor: |
| if isinstance(token_ids, torch.Tensor): |
| row = token_ids.detach().cpu() |
| if row.dtype == torch.bool or row.dtype.is_floating_point or row.dtype.is_complex: |
| raise ValueError("token ids must be integers") |
| row = row.to(torch.long) |
| else: |
| values = np.asarray(token_ids) |
| if values.dtype.kind not in {"i", "u"} and values.size: |
| raise ValueError("token ids must be integers") |
| row = torch.from_numpy(np.asarray(values, dtype=np.int64)) |
| if row.ndim != 1: |
| raise ValueError(f"token id row must be 1D, got {tuple(row.shape)}") |
| return row |
|
|
| def _checked_decode_row(self, token_ids: TokenIds) -> torch.Tensor: |
| row = self._to_id_row(token_ids) |
| invalid = (row < 0) | (row >= int(self._action_vocab_size)) |
| if torch.any(invalid): |
| bad = row[invalid][:5].tolist() |
| raise ValueError(f"token ids {bad} out of range [0, {int(self._action_vocab_size)})") |
| return self._strip_bar_sentinels(row) |
|
|
| def decode_token_ids_to_actions( |
| self, |
| action_token_ids: TokenIds, |
| *, |
| time_horizon: int | None = None, |
| action_dim: int | None = None, |
| decode_kwargs: dict | None = None, |
| ) -> DecodeResult: |
| """Decode one action-token id row back to a ``(T, D)`` action chunk.""" |
| row = self._checked_decode_row(action_token_ids) |
| actions_np, absent_keys_list = self._decode_action_indices( |
| [row.tolist()], |
| time_horizon=time_horizon, |
| action_dim=action_dim, |
| decode_kwargs=decode_kwargs, |
| ) |
| return DecodeResult( |
| action=torch.from_numpy(actions_np)[0], |
| absent_keys=absent_keys_list[0] if absent_keys_list else set(), |
| ) |
|
|
| 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 a batch of id rows to actions ``(B, T, D)``.""" |
| batch_action_idx: ActionIdxBatch = [ |
| self._checked_decode_row(a).tolist() for a in action_token_ids |
| ] |
| actions_np, _ = self._decode_action_indices( |
| batch_action_idx, |
| time_horizon=time_horizon, |
| action_dim=action_dim, |
| decode_kwargs=decode_kwargs, |
| ) |
| return torch.from_numpy(actions_np) |
|
|
| @property |
| def vocab_size(self) -> int: |
| return int(self._action_vocab_size) |
|
|
| @abstractmethod |
| def _encode_action_indices( |
| self, action: torch.Tensor, encode_kwargs: dict | None = None |
| ) -> ActionIdxBatch: ... |
|
|
| @abstractmethod |
| 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]]]: ... |
|
|
