refactor: remove habitat_sim dependency in motor_policies.py (#79)

Author: tristanls

src/tbp/monty/frameworks/models/motor_policies.py

@@ -15,8 +15,8 @@
 import os
 from typing import Any, Callable, Dict, List, Tuple, Type, Union, cast
 
-import habitat_sim.utils as hab_utils
 import numpy as np
+import quaternion as qt
 import scipy.ndimage
 from scipy.spatial.transform import Rotation as rot  # noqa: N813
 
@@ -1150,7 +1150,7 @@ def tangential_direction(self) -> VectorXYZ:
             self.tangential_angle * (1 - self.alpha) + new_target_direction * self.alpha
         )
 
-        direction = hab_utils.quat_rotate_vector(
+        direction = qt.rotate_vectors(
             self.state[self.agent_id]["rotation"],
             [
                 np.cos(self.tangential_angle - np.pi / 2),
@@ -1220,19 +1220,14 @@ def get_inverse_agent_rot(self):
         identity pose, which will be aquired by transforming the original pose by the
         inverse
 
-        NB Magnum appears to be a library used by Habitat; i.e. it is still grounded
-        in quaternions
-
         Returns:
             Inverse quaternion rotation.
         """
-        inverse_magnum_rotation = hab_utils.common.quat_to_magnum(
-            self.state["agent_id_0"]["rotation"]
-        ).inverted()
-        inverse_quaternion_rotation = hab_utils.common.quat_from_magnum(
-            inverse_magnum_rotation
-        )
-        return inverse_quaternion_rotation
+        # Note that quaternion format is [w, x, y, z]
+        [w, x, y, z] = qt.as_float_array(self.state[self.agent_id]["rotation"])
+        # Note that scipy.spatial.transform.Rotation (v1.10.0) format is [x, y, z, w]
+        [x, y, z, w] = rot.from_quat([x, y, z, w]).inv().as_quat()
+        return qt.quaternion(w, x, y, z)
 
     def orienting_angle_from_normal(self, orienting: str) -> float:
         """Compute turn angle to face the object.
@@ -1250,7 +1245,7 @@ def orienting_angle_from_normal(self, orienting: str) -> float:
 
         inverse_quaternion_rotation = self.get_inverse_agent_rot()
 
-        rotated_point_normal = hab_utils.quat_rotate_vector(
+        rotated_point_normal = qt.rotate_vectors(
             inverse_quaternion_rotation, original_point_normal
         )
         x, y, z = rotated_point_normal
@@ -1603,9 +1598,7 @@ def perform_pc_guided_step(self):
         # Rotate the tangential vector to be in the coordinate frame of the sensory
         # agent (rather than the global reference frame of the environment)
         inverse_quaternion_rotation = self.get_inverse_agent_rot()
-        rotated_form = hab_utils.quat_rotate_vector(
-            inverse_quaternion_rotation, selected_pc_dir
-        )
+        rotated_form = qt.rotate_vectors(inverse_quaternion_rotation, selected_pc_dir)
 
         # Before updating the representation and removing z-axis direction, check
         # for movements defined in the z-axis
@@ -1654,7 +1647,7 @@ def perform_pc_guided_step(self):
             self.reset_pc_buffers()
             self.following_heading_counter = 0
 
-            return hab_utils.quat_rotate_vector(
+            return qt.rotate_vectors(
                 self.state["agent_id_0"]["rotation"],
                 self.tangential_vec,
             )
@@ -1675,7 +1668,7 @@ def perform_pc_guided_step(self):
         self.following_pc_counter += 1
         self.continuous_pc_steps += 1
 
-        return hab_utils.quat_rotate_vector(
+        return qt.rotate_vectors(
             self.state["agent_id_0"]["rotation"],
             self.tangential_vec,
         )
@@ -1739,7 +1732,7 @@ def perform_standard_tang_step(self):
 
         self.following_heading_counter += 1
 
-        return hab_utils.quat_rotate_vector(
+        return qt.rotate_vectors(
             self.state["agent_id_0"]["rotation"],
             self.tangential_vec,
         )
@@ -1895,10 +1888,9 @@ def avoid_revisiting_locations(
             # locations to the headings (also in the reference frame of the agent) that
             # we might take
             # TODO could vectorize this
-            rotated_locs = [
-                hab_utils.quat_rotate_vector(inverse_quaternion_rotation, point)
-                for point in adjusted_prev_locs
-            ]
+            rotated_locs = qt.rotate_vectors(
+                inverse_quaternion_rotation, adjusted_prev_locs
+            )
 
             # Until we have not found a direction that we can guarentee is
             # in a new heading, continue to attempt new directions