Method __call__

model/action_tokenizer.py:177–188

Discretize continuous actions to tokens. action: np.ndarray, (n, 3), continuous actions in Cartesian or Spherical coordinates. return: np.ndarray, (n,), tokens.

(self, action: np.ndarray)

        self.yaw_bins = np.array(bin_policy["yaw_bins"])

    def __call__(self, action: np.ndarray) -> List[str]:
        """Discretize continuous actions to tokens.
        action: np.ndarray, (n, 3), continuous actions in Cartesian or Spherical coordinates.
        return: np.ndarray, (n,), tokens.
        """
        roll, pitch, yaw = action[:, 0], action[:, 1], action[:, 2]
        disc_roll = np.clip(np.digitize(roll, self.roll_bins) - 1, 0, self.num_roll_bins - 1)
        disc_pitch = np.clip(np.digitize(pitch, self.pitch_bins) - 1, 0, self.num_pitch_bins - 1)
        disc_yaw = np.clip(np.digitize(yaw, self.yaw_bins) - 1, 0, self.num_yaw_bins - 1)

        ids = disc_roll * self.NP + disc_pitch * self.num_yaw_bins + disc_yaw
        return self.token_array[ids]

    def decode_token_ids_to_actions(self, action_token_id: Union[np.int64, np.ndarray]) -> np.ndarray:
        """decode token ids to continuous actions.