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)
| 103 | return x, y, z |
| 104 | |
| 105 | def __call__(self, action: np.ndarray) -> List[str]: |
| 106 | """Discretize continuous actions to tokens. |
| 107 | action: np.ndarray, (n, 3), continuous actions in Cartesian or Spherical coordinates. |
| 108 | return: np.ndarray, (n,), tokens. |
| 109 | """ |
| 110 | if self.use_spherical: |
| 111 | theta, phi, r = self.cartesian_to_spherical(action[:, 0], action[:, 1], action[:, 2]) |
| 112 | else: |
| 113 | theta, phi, r = action[:, 0], action[:, 1], action[:, 2] |
| 114 | |
| 115 | disc_theta = np.digitize(theta, self.theta_bins[1:-1]) # b |
| 116 | disc_phi = np.digitize(phi, self.phi_bins[1:-1]) |
| 117 | disc_r = np.digitize(r, self.r_bins[1:-1]) |
| 118 | ids = disc_theta * self.NP + disc_phi * self.num_r_bins + disc_r |
| 119 | return self.token_array[ids] |
| 120 | |
| 121 | def decode_token_ids_to_actions(self, action_token_id: np.ndarray) -> np.ndarray: |
| 122 | """decode token ids to continuous actions. |
nothing calls this directly
no test coverage detected