"""Action tokenizer mixin for loading tokenizers and action mappings.""" from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import torch from transformers import AutoProcessor from wall_x.utils.constant import is_action_dataset_name # Delay imports so missing optional packages do not fail at import time try: from spatial_tokenizer.spatial_tokenizer import SpatialActionTokenizer except ImportError: SpatialActionTokenizer = None class ActionTokenizerMixin(ABC): """Base class for action tokenizers""" def __init__(self): self._tokenizer = None self.action_normalizer = None self._tokenizer_type: str = "" self._action_mapper_cache: Optional[Dict] = None # Cached action_mapper self.dllm = False self.input_placeholder_flag = False @property def tokenizer_type(self) -> str: """Return the tokenizer type identifier""" return self._tokenizer_type @property def tokenizer(self): """Return the underlying tokenizer instance""" return self._tokenizer @property def action_mapper(self) -> Optional[Dict]: """Return the cached action_mapper""" return self._action_mapper_cache @abstractmethod def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: """ Load the tokenizer instance Args: config: Configuration dictionary normalizer: Normalizer device: Device, usually "cpu" for training and "cuda" for inference Returns: tokenizer instance """ pass @abstractmethod def get_val_tokenizer(self, config: dict) -> Any: """ Get the tokenizer used for validation/inference Args: config: Configuration dictionary Returns: validation tokenizer instance """ pass @abstractmethod def get_special_tokens(self) -> List[str]: """ Return special tokens to add to the vocabulary Returns: list of token strings """ pass def get_all_special_tokens(self) -> Tuple[List[str], Optional[List[str]]]: """ Return all special tokens and keep <|action_token_0|> before AR action tokens Returns: (new_tokens, special_tokens) tuple """ tokens, special_tokens = self.get_special_tokens() if "<|action_token_0|>" not in tokens: tokens.insert(0, "<|action_token_0|>") return tokens, special_tokens @abstractmethod def build_action_mapper(self, processor) -> Optional[Dict]: """ Build action_mapper Args: processor: HuggingFace processor,used to convert tokens to IDs Returns: action_mapper dictionary; format depends on tokenizer type """ pass @abstractmethod def get_action_token_list(self, processor) -> List[int]: """ Get the action token ID list Args: processor: HuggingFace processor Returns: action token ID list """ pass @abstractmethod def decode_action( self, output_ids: torch.Tensor, action_mapper: Dict, action_horizon: int, action_dim: int, device: torch.device, proprioception: Optional[torch.Tensor] = None, dof_mask: Optional[torch.Tensor] = None, robot_type_id: Optional[int] = None, state: Optional[torch.Tensor] = None, ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: """ Unified decoding interface Args: output_ids: model output token IDs [1, seq_len] action_mapper: action_mapper dictionary action_horizon: action horizon action_dim: action dimension device: Device proprioception: proprioception, normalized when required by a tokenizer dof_mask: DOF mask when required by a tokenizer robot_type_id: robot type ID when required by a tokenizer state: state when required by fast/spatial tokenizers Returns: (predict_action, decode_success) - predict_action: decoded action [T, action_dim] or None - decode_success: whether decoding succeeded """ pass @abstractmethod def compute_accuracy( self, logits: torch.Tensor, labels: torch.Tensor, action_mapper: Dict, action_token_id_set: Dict, ) -> Dict[str, torch.Tensor]: """ Compute accuracy metrics Args: logits: model output logits labels: Labels action_mapper: action_mapper dictionary action_token_id_set: set of action token IDs Returns: accuracy metric dictionary, such as {"action_accuracy": tensor, ...} """ pass @abstractmethod def get_accuracy_keys(self) -> List[str]: """ Return accuracy metric keys for logging Returns: list of key strings """ pass @property @abstractmethod def vocab_size(self) -> int: """Return the vocabulary size""" pass @property @abstractmethod def uses_dof_mask_for_unnorm(self) -> bool: """ Whether unnormalization requires dof_mask Returns: True: dof_mask is required by fast/spatial tokenizers False: dof_mask is not required and full dimensions are returned """ pass @property @abstractmethod def needs_action_crop(self) -> bool: """ Whether actions must be clipped before encoding Returns: True: clip by chunk_size and dof_mask for fast/spatial tokenizers False: no clipping; the encoder handles the full sequence internally """ pass @abstractmethod def encode_to_tokens( self, actions: torch.Tensor, obs_state: Optional[torch.Tensor] = None, dof_mask: Optional[torch.Tensor] = None, robot_type_ids: Optional[List[int]] = None, is_train: bool = True, ) -> List[List[str]]: """ Encode actions into token strings for training data processing Args: actions: Normalized actions - fast/spatial: List[Tensor], each [T, D] (clipped) - v3.1 delta: Tensor [B, T, D] (full sequence) obs_state: observation state when required by a tokenizer[B, obs_horizon, D] dof_mask: DOF mask[B, T, D] robot_type_ids: robot type ID list when required by a tokenizer Returns: List[List[str]]: token string list for each sample """ pass def init_inference(self, robot_type: Optional[str] = None) -> None: """ Inference initialization hook; subclasses may override Args: robot_type: robot type name """ pass def prepare_action_for_ar_encoding( self, ar_actionchunk: torch.Tensor, dataset_names: List[str], agent_pos: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Prepare action data before AR encoding; subclasses may override SpatialVLA needs waypoint selection; other tokenizers return normalized_action directly Args: ar_actionchunk: [B, T, D] actions before normalization dataset_names: [str] dataset name list agent_pos: [B, T, D] agent positions required by SpatialVLA Returns: prepared action data """ if self.action_normalizer is not None: ar_actionchunk = self.action_normalizer.normalize_data( ar_actionchunk, dataset_names ) return ar_actionchunk def get_robot_type_ids( self, uids: List[Optional[str]], dataset_names: List[str], ) -> Optional[List[int]]: """ Get robot_type_ids; subclasses may override Only the v3.1 delta tokenizer needs this; other tokenizers return None Args: uids: UID list dataset_names: dataset name list Returns: robot_type_id list or None """ return None def modify_inputs_for_dllm( self, inputs: Dict[str, torch.Tensor], processor, sample_time: torch.Tensor, dataset_names: List[str], ): return NotImplementedError class FastTokenizerMixin(ActionTokenizerMixin): """Fast tokenizer implementation""" def __init__(self): super().__init__() self._tokenizer_type = "fast" def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: """Load the fast tokenizer""" self._tokenizer = AutoProcessor.from_pretrained( config["action_tokenizer_path"], trust_remote_code=True ) self.action_normalizer = normalizer return self._tokenizer def get_val_tokenizer(self, config: dict) -> Any: return self._tokenizer def get_special_tokens(self) -> List[str]: """Return special tokens for the fast tokenizer""" if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") # tokens = ["<|ar_action|>", "<|ar_pad|>"] #temporary compatibility with existing checkpoints tokens = [] for i in range(self._tokenizer.vocab_size): tokens.append(f"<|action_token_{i}|>") return tokens, None def build_action_mapper(self, processor) -> Dict[int, int]: """ Build the cached action_mapper for the fast tokenizer Returns: Dict[token_id, action_idx] """ # Return cached value if available if self._action_mapper_cache is not None: return self._action_mapper_cache if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") action_mapper = {} for i in range(self._tokenizer.vocab_size): token = f"<|action_token_{i}|>" token_id = processor.tokenizer.convert_tokens_to_ids(token) action_mapper[token_id] = i self._action_mapper_cache = action_mapper return action_mapper def get_action_token_list(self, processor) -> List[int]: """Get the action token ID list""" if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") action_token_list = [] for i in range(self._tokenizer.vocab_size): token_id = processor.tokenizer.convert_tokens_to_ids( f"<|action_token_{i}|>" ) action_token_list.append(token_id) return action_token_list def decode_action( self, output_ids: torch.Tensor, action_mapper: Dict, action_horizon: int, action_dim: int, device: torch.device, proprioception: Optional[torch.Tensor] = None, dof_mask: Optional[torch.Tensor] = None, robot_type_id: Optional[int] = None, state: Optional[torch.Tensor] = None, ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: """Fast tokenizer decoding""" action_id = [] for token_id_i in output_ids[0]: if token_id_i.item() in action_mapper: action_id.append(action_mapper[token_id_i.item()]) if len(action_id) == 0: return np.zeros((action_horizon, action_dim)), False predict_action = self._tokenizer.decode( [action_id], time_horizon=action_horizon, action_dim=action_dim ) # Check whether decoding succeeded decode_success = False if isinstance(predict_action, np.ndarray): decode_success = np.sum(predict_action) != 0 elif isinstance(predict_action, torch.Tensor): decode_success = predict_action.sum().item() != 0 return predict_action, decode_success def compute_accuracy( self, logits: torch.Tensor, labels: torch.Tensor, action_mapper: Dict, action_token_id_set: Dict, ) -> Dict[str, torch.Tensor]: """Compute fast tokenizer accuracy""" result = {} if len(action_token_id_set.get("action_token_list", [])) > 0: shift_logits = logits[..., :-1, :].contiguous() action_preds = shift_logits.argmax(dim=-1) shift_labels = labels[..., 1:].contiguous() action_mask = shift_labels > action_token_id_set["action_token_list"][0] correct_preds = (action_preds == shift_labels) & action_mask action_accuracy = correct_preds.sum().float() / action_mask.sum().float() result["action_accuracy"] = action_accuracy return result def get_accuracy_keys(self) -> List[str]: """Return fast tokenizer accuracy keys""" return ["action_accuracy"] @property def vocab_size(self) -> int: if self._tokenizer is None: return 0 return self._tokenizer.vocab_size @property def uses_dof_mask_for_unnorm(self) -> bool: """fast tokenizer requires dof_mask""" return True @property def needs_action_crop(self) -> bool: return True @property def inference_ar_steps_for_dllm(self) -> int: return self.max_length def encode_to_tokens( self, actions: List, obs_state: Optional[torch.Tensor] = None, dof_mask: Optional[torch.Tensor] = None, robot_type_ids: Optional[List[int]] = None, is_train: bool = True, ) -> List[List[str]]: """ Fast tokenizer encoding Args: actions: List[Tensor/ndarray], each [T, D] (clipped) """ if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") all_action_tokens = [] for i in range(len(actions)): action = actions[i] if isinstance(action, torch.Tensor): action = action.cpu().numpy() token_id = self._tokenizer(action) action_tokens = [f"<|action_token_{idx}|>" for idx in token_id[0]] all_action_tokens.append(action_tokens) return all_action_tokens def modify_inputs_for_dllm( self, inputs: Dict[str, torch.Tensor], processor, sample_time: torch.Tensor, dataset_names: List[str], ): # Untested input_ids = inputs["input_ids"] labels = inputs["labels"] prefix_length = inputs["prefix_length"] bs, seqlen = input_ids.shape ar_token_length = self.max_length ar_step_num = ar_token_length device = input_ids.device dtype = input_ids.dtype placeholder_ids = torch.tensor( processor.placeholder_seq, device=device, dtype=dtype ) if not torch.is_tensor(sample_time): sample_time = torch.tensor(sample_time, device=device, dtype=torch.float32) else: sample_time = sample_time.to(device=device, dtype=torch.float32) sample_time = sample_time.clamp(0.0, 1.0) noisy_steps_per_sample = ( torch.ceil((1.0 - sample_time) * (ar_step_num + 1)).long() - 1 ) noisy_steps_per_sample = noisy_steps_per_sample.clamp(min=0, max=ar_step_num) start = prefix_length - ar_token_length - 2 end = prefix_length - 2 # Prepare the noise sequence for each sample noise_seqs = placeholder_ids.repeat(bs, 1) ar_len = end - start if noise_seqs.size(1) != ar_len: # These should usually match; defensively truncate to ar_len noise_seqs = noise_seqs[:, :ar_len] rand = torch.rand(bs, ar_len, device=device) perm = rand.argsort(dim=-1) ranks = perm.argsort(dim=-1) noisy_mask = ranks < noisy_steps_per_sample.view(-1, 1) ar_input = input_ids[:, start:end] ar_labels = labels[:, start + 1 : end + 1] ar_input[noisy_mask] = noise_seqs[noisy_mask] ar_labels[~noisy_mask] = -100 inputs["input_ids"] = input_ids inputs["labels"] = labels return inputs def update_placeholder_mask(self, processor, prefix_length, input_ids): # Untested ar_action_mask = torch.zeros_like(input_ids) inc = torch.arange( 1, self.max_length + 1, ) inc = inc.unsqueeze(0).expand(ar_action_mask.size(0), -1) # [bs, ar_len] ar_action_mask[ :, prefix_length - self.max_length - 2 : prefix_length - 2, ] = inc ar_action_mask[:, prefix_length - 2 : prefix_length] = -1 # eos return {"ar_action_mask": ar_action_mask} def get_placeholder_for_dllm(self): placeholder_seq = ["<|ar_action|>"] * self.max_length return placeholder_seq class SpatialVLATokenizerMixin(ActionTokenizerMixin): """SpatialVLA tokenizer implementation""" def __init__(self): super().__init__() self._tokenizer_type = "spatialvla" def load_tokenizer(self, config: dict, normalizer, device: str = "cpu") -> Any: """Load the SpatialVLA tokenizer""" if SpatialActionTokenizer is None: raise ImportError( "SpatialActionTokenizer is not installed. " "Please install spatial_tokenizer package." ) self._tokenizer = SpatialActionTokenizer( normalizer=normalizer, augment_ratio=config.get("augment_ratio", 0.0), max_waypoints=config.get("max_waypoints", 5), with_gripper=config.get("with_gripper", True), single_arm=config.get("single_arm", False), ) self._val_tokenizer = None self.config = config self.dllm = config.get("dllm", False) self.input_placeholder_flag = config.get("input_placeholder_flag", False) self.action_normalizer = normalizer self.with_gripper = config.get("with_gripper", True) return self._tokenizer def get_placeholder_for_dllm(self): if self.with_gripper: placeholder_seq = [ "<|left_xyz|>", "<|left_rpy|>", "<|left_gripper|>", "<|right_xyz|>", "<|right_rpy|>", "<|right_gripper|>", ] else: placeholder_seq = [ "<|left_xyz|>", "<|left_rpy|>", "<|right_xyz|>", "<|right_rpy|>", ] if self._tokenizer.single_arm: placeholder_seq = placeholder_seq[len(placeholder_seq) // 2 :] placeholder_seq = placeholder_seq * self._tokenizer.max_waypoints return placeholder_seq def get_val_tokenizer(self, config: dict) -> Any: if self._val_tokenizer: return self._val_tokenizer self._val_tokenizer = SpatialActionTokenizer( normalizer=self.action_normalizer, augment_ratio=0, max_waypoints=self.config.get("max_waypoints", 5), with_gripper=self.config.get("with_gripper", True), single_arm=self.config.get("single_arm", False), ) return self._val_tokenizer def get_special_tokens(self) -> List[str]: """Return special tokens for the SpatialVLA tokenizer""" if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") tokens = [ "<|step|>", "<|left|>", "<|right|>", "<|move|>", ] # only for compatibility with existing checkpoints if self.input_placeholder_flag: special_tokens = [ "<|left_xyz|>", "<|left_rpy|>", "<|left_gripper|>", "<|right_xyz|>", "<|right_rpy|>", "<|right_gripper|>", ] tokens += special_tokens if not self._tokenizer.with_gripper: indices = [0, 1, 3, 4] special_tokens = [special_tokens[i] for i in indices] if self._tokenizer.single_arm: special_tokens = special_tokens[len(special_tokens) // 2 :] for i in range(self._tokenizer.vocab_size): tokens.append(f"<|action_token_{i}|>") return tokens, special_tokens def build_action_mapper(self, processor) -> Dict[int, int]: """ Build the cached action_mapper for the SpatialVLA tokenizer Returns: Dict[token_id, action_idx] """ # Return cached value if available if self._action_mapper_cache is not None: return self._action_mapper_cache if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") action_mapper = {} for i in range(self._tokenizer.vocab_size): token = f"<|action_token_{i}|>" token_id = processor.tokenizer.convert_tokens_to_ids(token) action_mapper[token_id] = i self._action_mapper_cache = action_mapper return action_mapper def get_action_token_list(self, processor) -> List[int]: """Get the action token ID list""" if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") action_token_list = [] for i in range(self._tokenizer.vocab_size): token_id = processor.tokenizer.convert_tokens_to_ids( f"<|action_token_{i}|>" ) action_token_list.append(token_id) return action_token_list def decode_action( self, output_ids: torch.Tensor, action_mapper: Dict, action_horizon: int, action_dim: int, device: torch.device, proprioception: Optional[torch.Tensor] = None, dof_mask: Optional[torch.Tensor] = None, robot_type_id: Optional[int] = None, state: Optional[torch.Tensor] = None, ) -> Tuple[Optional[Union[np.ndarray, torch.Tensor]], bool]: """SpatialVLA tokenizer decoding""" action_id = [] for token_id_i in output_ids[0]: if token_id_i.item() in action_mapper: action_id.append(action_mapper[token_id_i.item()]) if len(action_id) == 0: return np.zeros((action_horizon, action_dim)), False if state is not None: predict_action = self._tokenizer.decode( [action_id], state=state[0, 0, :action_dim], time_horizon=action_horizon, action_dim=action_dim, ) else: predict_action = self._tokenizer.decode( [action_id], time_horizon=action_horizon, action_dim=action_dim ) # Check whether decoding succeeded decode_success = False if isinstance(predict_action, np.ndarray): decode_success = np.sum(predict_action) != 0 elif isinstance(predict_action, torch.Tensor): decode_success = predict_action.sum().item() != 0 return predict_action, decode_success def compute_accuracy( self, logits: torch.Tensor, labels: torch.Tensor, action_mapper: Dict, action_token_id_set: Dict, ) -> Dict[str, torch.Tensor]: """Compute SpatialVLA tokenizer accuracy, same as fast""" result = {} if len(action_token_id_set.get("action_token_list", [])) > 0: shift_logits = logits[..., :-1, :].contiguous() action_preds = shift_logits.argmax(dim=-1) shift_labels = labels[..., 1:].contiguous() action_mask = shift_labels > action_token_id_set["action_token_list"][0] correct_preds = (action_preds == shift_labels) & action_mask action_accuracy = correct_preds.sum().float() / action_mask.sum().float() result["action_accuracy"] = action_accuracy return result def get_accuracy_keys(self) -> List[str]: """Return SpatialVLA tokenizer accuracy keys""" return ["action_accuracy"] @property def vocab_size(self) -> int: if self._tokenizer is None: return 0 return self._tokenizer.vocab_size @property def uses_dof_mask_for_unnorm(self) -> bool: """SpatialVLA tokenizer requires dof_mask""" return True @property def needs_action_crop(self) -> bool: return False @property def inference_ar_steps_for_dllm(self) -> int: return self._tokenizer.max_waypoints def prepare_action_for_ar_encoding( self, ar_actionchunk: torch.Tensor, dataset_names: List[str], agent_pos: Optional[torch.Tensor] = None, ) -> List[torch.Tensor]: step = ar_actionchunk.shape[1] // self._tokenizer.max_waypoints indices = np.arange(0, ar_actionchunk.shape[1], step)[ : self._tokenizer.max_waypoints ] ar_action = ar_actionchunk[:, indices, :] ar_action = self.action_normalizer.normalize_data(ar_action, dataset_names) return ar_action def encode_to_tokens( self, actions: List, obs_state: Optional[torch.Tensor] = None, dof_mask: Optional[torch.Tensor] = None, robot_type_ids: Optional[List[int]] = None, is_train: bool = True, ) -> List[List[str]]: """ SpatialVLA tokenizer encoding Args: actions: List[Tensor/ndarray], each [T, D] (clipped) """ if self._tokenizer is None: raise RuntimeError("Tokenizer not loaded. Call load_tokenizer first.") if isinstance(actions, torch.Tensor): # Convert torch tensors to numpy arrays actions = actions.cpu().numpy() tokenizer = self._tokenizer if is_train else self.get_val_tokenizer({}) token_ids_group = tokenizer.batch_encode(actions) all_action_tokens = [] for i in range(len(token_ids_group)): token_ids = np.array(token_ids_group[i]).reshape(-1) action_token = [f"<|action_token_{i}|>" for i in token_ids] all_action_tokens.append(action_token) return all_action_tokens def modify_inputs_for_dllm( self, inputs: Dict[str, torch.Tensor], processor, sample_time: torch.Tensor, dataset_names: List[str], ): """ Prepare AR DLLM inputs by encoding sample_time into input_ids and labels Noise injection strategy: - Split [0, 1] into N equal parts (N = ar_step_num) - For each sample, compute the number of noisy steps k in [1, N] k = clamp(ceil((1 - t) * N), 1, N) -> smaller t means more noise; values closer to 1 mean less noise, with at least one noised step - seq is the noise sequence formed by concatenating ar_step_num placeholder_seq blocks Each step maps to len(placeholder_seq) consecutive tokens - Noise injection overwrites matching tokens in input_ids / labels step by step from seq """ input_ids = inputs["input_ids"] labels = inputs["labels"] prefix_length = inputs["prefix_length"] # int or tensor(1,) bs, seqlen = input_ids.shape # Total number of AR steps ar_step_num = self._tokenizer.max_waypoints # N step_token_len = len(processor.placeholder_seq) # Number of tokens per step ar_token_length = ar_step_num * step_token_len # Total AR token length device = input_ids.device dtype = input_ids.dtype placeholder_ids = torch.tensor( processor.placeholder_seq, device=device, dtype=dtype ) # (step_token_len,) noise_seq = placeholder_ids.repeat(ar_step_num) # (ar_token_length,) if not torch.is_tensor(sample_time): sample_time = torch.tensor(sample_time, device=device, dtype=torch.float32) else: sample_time = sample_time.to(device=device, dtype=torch.float32) sample_time = sample_time.clamp(0.0, 1.0) noisy_steps_per_sample = ( torch.ceil((1.0 - sample_time) * (ar_step_num + 1)).long() - 1 ) noisy_steps_per_sample = noisy_steps_per_sample.clamp(min=0, max=ar_step_num) start = prefix_length - ar_token_length - 2 end = prefix_length - 2 action_idx = 0 # action sample counter for i in range(bs): if not is_action_dataset_name(dataset_names[i]): continue k = noisy_steps_per_sample[action_idx].item() action_idx += 1 perm = torch.randperm(ar_step_num, device=device) chosen_steps = perm[:k] # (k,) step_mask = torch.zeros(ar_step_num, dtype=torch.bool, device=device) step_mask[chosen_steps] = True token_mask = step_mask.repeat_interleave( step_token_len ) # (ar_token_length,) ignore_mask = ~token_mask # Take a view of the current sample's AR span for mask-based replacement cur_input_view = input_ids[i, start:end] cur_label_view = labels[ i, start + 1 : end + 1 ] # shift placeholders one position to the right # Overwrite selected tokens with noise cur_input_view[token_mask] = noise_seq[token_mask] cur_label_view[ignore_mask] = -100 inputs["input_ids"] = input_ids inputs["labels"] = labels return inputs def update_positional_masks_for_dllm( self, positional_masks, inputs, processor, visible_predict_ar_ratio=1 ): if self.dllm and self.input_placeholder_flag: mask = self.update_placeholder_mask( processor, inputs["prefix_length"], inputs["input_ids"], ) positional_masks.update(mask) if ( np.random.rand() < visible_predict_ar_ratio ): # FIXME ar_visible is temporarily decided per batch; mixed settings are untested and need optimization. positional_masks["ar_visible"] = True if "ar_predict_token_positions" in positional_masks: del positional_masks["ar_predict_token_positions"] # positional_masks["ar_predict_token_positions"] = None else: positional_masks["ar_visible"] = False return positional_masks def update_placeholder_mask(self, processor, prefix_length, input_ids): ar_action_mask = torch.zeros_like(input_ids) current_index = 1 step_len = len(processor.placeholder_seq) for b, seq in enumerate(input_ids): for i in range(self._tokenizer.max_waypoints): start_idx = ( prefix_length - (self._tokenizer.max_waypoints - i) * step_len - 2 ) end_idx = start_idx + step_len ar_action_mask[b, start_idx:end_idx] = current_index current_index += 1 ar_action_mask[b, end_idx:prefix_length] = -1 # eos return {"ar_action_mask": ar_action_mask} # ============================================================ # Factory function # ============================================================ _TOKENIZER_REGISTRY: Dict[str, type] = { "fast": FastTokenizerMixin, "spatialvla": SpatialVLATokenizerMixin, } def get_action_tokenizer_mixin(tokenizer_type: str) -> ActionTokenizerMixin: """ Get the mixin instance for a tokenizer type Args: tokenizer_type: tokenizer type, supports "fast", "spatialvla" Returns: ActionTokenizerMixin instance Raises: ValueError: Unsupported tokenizer type """ if tokenizer_type not in _TOKENIZER_REGISTRY: raise ValueError( f"Unsupported action tokenizer type: {tokenizer_type}. " f"Supported types: {list(_TOKENIZER_REGISTRY.keys())}" ) return _TOKENIZER_REGISTRY[tokenizer_type]()