Method decode_actions

model/processing_spatialvla.py:216–254

(
        self,
        generation_outputs: torch.Tensor,
        unnorm_key: Optional[str] = None,
    )

        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

    def decode_actions(
        self,
        generation_outputs: torch.Tensor,
        unnorm_key: Optional[str] = None,
    ) -> Dict[str, torch.Tensor]:
        action_token_num = 3  # translation + rotation + gripper
        predicted_action_token_ids = generation_outputs[0, : action_token_num * self.action_chunk_size].detach().cpu().long().numpy()
        assert self.tokenizer.eos_token != predicted_action_token_ids[-1], "[error] actions contain EOS token, please check you truncation settings!"

        if predicted_action_token_ids.shape[0] < action_token_num * self.action_chunk_size:  # pad with zeros
            logger.warning(f"Padding zero action!")
            predicted_action_token_ids = np.concatenate(
                [
                    predicted_action_token_ids,
                    np.zeros(action_token_num * self.action_chunk_size - predicted_action_token_ids.shape[0], dtype=np.longlong),
                ]
            )
        predicted_action_token_ids = predicted_action_token_ids.reshape(-1, action_token_num)
        normalized_action_chunks = self.action_tokenizer.decode_token_ids_to_actions(predicted_action_token_ids)

        if unnorm_key is None:
            logger.warning(f"unnorm_key {unnorm_key} is not in statistics, use next one")
            unnorm_key = next(self.statistics.keys())
        action_norm_stats = self.statistics[unnorm_key]["action"]

        action_dim = len(action_norm_stats["q01"])
        mask = np.array(action_norm_stats.get("mask", np.ones(action_dim)), dtype=bool)
        action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])

        actions = []
        for normalized_actions in normalized_action_chunks:
            action = np.where(
                mask,
                0.5 * (normalized_actions + 1) * (action_high - action_low) + action_low,
                normalized_actions,
            )
            actions.append(action)
        actions = np.stack(actions)
        return {"actions": actions, "action_ids": predicted_action_token_ids}