Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

HITL - Add missing sim type hints. #1989

Merged
merged 4 commits into from
Jun 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 10 additions & 5 deletions examples/hitl/rearrange_v2/world.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand All @@ -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]
Expand All @@ -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
Expand Down
7 changes: 6 additions & 1 deletion habitat-hitl/habitat_hitl/core/gui_input.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand All @@ -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)
Expand Down
12 changes: 9 additions & 3 deletions habitat-hitl/habitat_hitl/environment/gui_placement_helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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."""
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
17 changes: 14 additions & 3 deletions habitat-hitl/habitat_hitl/environment/hablab_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Expand Down