Class ActionTokenizer

model/action_tokenizer.py:14–57

ACTION_TOKEN = '<ACTION{:05d}>'

class ActionTokenizer:
    def __init__(
        self,
        tokenizer: PreTrainedTokenizerBase,
        num_bins: int = 256,
        min_action: int = -1,
        max_action: int = 1,
    ):
        self._vocab_size = num_bins
        self.tokenizer = tokenizer
        self.min_action, self.max_action = min_action, max_action
        self.bin_centers = np.linspace(min_action, max_action, num_bins)

        # add special action tokens to language tokenizer
        token_list = [ACTION_TOKEN.format(i) for i in range(self._vocab_size)]
        self.token_array = np.array(token_list)
        
        num_new_tokens = self.tokenizer.add_tokens(token_list, special_tokens=True)
        print(f"Add {num_new_tokens} TRANSLATION TOKENS, tokenizer vocab size {self.tokenizer.vocab_size} / {len(tokenizer)}")

        self.action_token_begin_idx = self.token_start_idx = self.tokenizer.convert_tokens_to_ids(self.token_array[0])
        self.token_end_idx = self.tokenizer.convert_tokens_to_ids(self.token_array[-1])

    def __call__(self, action: np.ndarray) -> List[str]:
        """Discretize continuous actions to tokens.
        action: np.ndarray, (n, 7), continuous actions in Cartesian or Spherical coordinates.
        return: np.ndarray, (n, 7), tokens.
        """
        action = np.clip(action, a_min=float(self.min_action), a_max=float(self.max_action))
        ids = np.digitize(action, self.bin_centers, right=True)  # [0, 255]
        return self.token_array[ids]

    def decode_token_ids_to_actions(self, action_token_id: np.ndarray) -> np.ndarray:
        """decode token ids to continuous actions.
        action_token_id: np.ndarray, (n, 7), token ids.
        return: np.ndarray, (n, 7), continuous actions
        """
        ids = action_token_id - self.action_token_begin_idx
        ids = np.clip(ids, a_min=0, a_max=self._vocab_size - 1)
        return self.bin_centers[ids]

    @property
    def vocab_size(self) -> int:
        return self._vocab_size

class TranslationTokenizer:
    def __init__(