HITL - Add missing sim type hints. (facebookresearch#1989)

* Add sim types to world.py.

* Add type hints to gui_placement_helper.

* Add type hints to hablab_utils.py.

* Add mouse ray type.

examples/hitl/rearrange_v2/world.py

@@ -11,7 +11,10 @@
     ArticulatedAgentManager,
 )
 from habitat.tasks.rearrange.rearrange_sim import RearrangeSim
-from habitat_sim.physics import ManagedArticulatedObject
+from habitat_sim.physics import (
+    ManagedBulletArticulatedObject,
+    ManagedBulletRigidObject,
+)
 
 
 class World:
@@ -68,7 +71,7 @@ def reset(self) -> None:
         self._interactable_object_ids = set()
         aom = sim.get_articulated_object_manager()
         all_ao: List[
-            ManagedArticulatedObject
+            ManagedBulletArticulatedObject
         ] = aom.get_objects_by_handle_substring().values()
         # Classify all non-root links.
         for ao in all_ao:
@@ -85,7 +88,9 @@ def reset(self) -> None:
                         if link_index:
                             sim_utilities.close_link(ao, link_index)
 
-    def get_rigid_object(self, object_id: int) -> Optional[Any]:
+    def get_rigid_object(
+        self, object_id: int
+    ) -> Optional[ManagedBulletRigidObject]:
         """Get the rigid object with the specified ID. Returns None if unsuccessful."""
         rom = self._sim.get_rigid_object_manager()
         return rom.get_object_by_id(object_id)
@@ -102,7 +107,7 @@ def get_link_index(self, object_id: int) -> int:
         )
         if (
             obj is not None
-            and isinstance(obj, ManagedArticulatedObject)
+            and isinstance(obj, ManagedBulletArticulatedObject)
             and object_id in obj.link_object_ids
         ):
             return obj.link_object_ids[object_id]
@@ -112,7 +117,7 @@ def get_agent_object_ids(self, agent_index: int) -> Set[int]:
         """Get the IDs of objects composing an agent (including links)."""
         # TODO: Cache
         sim = self._sim
-        agent_manager: ArticulatedAgentManager = sim.agents_mgr
+        agent_manager = sim.agents_mgr
         agent_object_ids: Set[int] = set()
         agent = agent_manager[agent_index]
         agent_ao = agent.articulated_agent.sim_obj

habitat-hitl/habitat_hitl/core/gui_input.py

@@ -6,8 +6,13 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Optional
+
 from habitat_hitl.core.key_mapping import KeyCode, MouseButton
 
+if TYPE_CHECKING:
+    from habitat_sim.geo import Ray
+
 
 class GuiInput:
     """
@@ -30,7 +35,7 @@ def __init__(self):
         self._mouse_button_up = set()
         self._relative_mouse_position = [0, 0]
         self._mouse_scroll_offset = 0.0
-        self._mouse_ray = None
+        self._mouse_ray: Optional[Ray] = None
 
     def validate_key(key):
         assert isinstance(key, KeyCode)

habitat-hitl/habitat_hitl/environment/gui_placement_helper.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import math
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 import magnum as mn
 
@@ -21,6 +21,10 @@
 FAR_AWAY_HIDDEN_POSITION = mn.Vector3(0, -1000, 0)
 DEFAULT_GRAVITY = mn.Vector3(0, -1, 0)
 
+if TYPE_CHECKING:
+    from habitat_sim.geo import Ray
+    from habitat_sim.physics import ManagedBulletRigidObject
+
 
 class GuiPlacementHelper:
     """Helper for placing objects from the GUI."""
@@ -35,7 +39,9 @@ def __init__(
         self._user_index = user_index
         self._gravity_dir = gravity_dir
 
-    def _snap_or_hide_object(self, ray, query_obj) -> tuple[bool, mn.Vector3]:
+    def _snap_or_hide_object(
+        self, ray: Ray, query_obj: ManagedBulletRigidObject
+    ) -> tuple[bool, mn.Vector3]:
         sim = self._app_service.sim
 
         assert query_obj.collidable
@@ -93,7 +99,7 @@ def _snap_or_hide_object(self, ray, query_obj) -> tuple[bool, mn.Vector3]:
 
         return True, adjusted_hit_pos
 
-    def update(self, ray, query_obj_id):
+    def update(self, ray: Ray, query_obj_id: int):
         sim = self._app_service.sim
         query_obj = sim.get_rigid_object_manager().get_object_by_id(
             query_obj_id

habitat-hitl/habitat_hitl/environment/hablab_utils.py

@@ -7,19 +7,30 @@
 # Utilities built on top of Habitat-lab
 
 
-def get_agent_art_obj(sim, agent_idx):
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from habitat.tasks.rearrange.rearrange_sim import RearrangeSim
+    from habitat_sim.physics import ManagedBulletArticulatedObject
+
+
+def get_agent_art_obj(
+    sim: "RearrangeSim", agent_idx: int
+) -> "ManagedBulletArticulatedObject":
     assert agent_idx is not None
     art_obj = sim.agents_mgr[agent_idx].articulated_agent.sim_obj
     return art_obj
 
 
-def get_agent_art_obj_transform(sim, agent_idx):
+def get_agent_art_obj_transform(sim: "RearrangeSim", agent_idx: int):
     assert agent_idx is not None
     art_obj = sim.agents_mgr[agent_idx].articulated_agent.sim_obj
     return art_obj.transformation
 
 
-def get_grasped_objects_idxs(sim, agent_idx_to_skip=None):
+def get_grasped_objects_idxs(
+    sim: "RearrangeSim", agent_idx_to_skip: Optional[int] = None
+):
     agents_mgr = sim.agents_mgr
 
     grasped_objects_idxs = []