| """ActionCodec2-only tensor and token-row validation helpers.""" |
|
|
| from __future__ import annotations |
|
|
| from collections.abc import Sequence |
|
|
| import numpy as np |
| import torch |
|
|
| from ._types import ActionIdxBatch, TokenIds |
|
|
|
|
| class BaseActionTokenizer: |
| """Shared action tensor methods, without policy-specific sentinels.""" |
|
|
| @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>``. |
| |
| This is exactly the text a VLM would see after registering |
| :attr:`new_action_tokens` — useful for eyeballing sequences and for |
| prompt-template debugging. |
| |
| Args: |
| token_ids: One row of action-token ids in ``[0, vocab_size)``. |
| |
| Returns: |
| list[str]: Token strings in row order. |
| """ |
| 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 _check_encode_kwargs(self, encode_kwargs: dict | None) -> dict: |
| """Normalize ``encode_kwargs`` to a dict, rejecting keys outside :attr:`encode_kwargs`.""" |
| 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}; it declares {sorted(self.encode_kwargs)}" |
| ) |
| return kw |
|
|
| @staticmethod |
| def _to_id_row(token_ids: TokenIds) -> torch.Tensor: |
| """Validate and normalize one integer id row to a 1-D long 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"}: |
| 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: |
| """Return one validated row without repairing invalid model output.""" |
| 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 row |
|
|
| 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)``. |
| |
| The decoder applies the configured action-space layout to each row. |
| """ |
| 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) |
|
|
