From 963fef02ea885e90cacecfd4a6ac882ae46d68c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mika=C3=ABl=20Dallaire=20C=C3=B4t=C3=A9?= <110583667+0mdc@users.noreply.github.com> Date: Sat, 10 Feb 2024 14:29:51 -0500 Subject: [PATCH] HITL: Add typing annotations to AppService (#1802) * Add type annotations to app_service. * Propagate AppService type annotations to the apps. * Move displaced comment. * Force keyword argument usage in AppService constructior. --- examples/hitl/basic_viewer/basic_viewer.py | 3 +- examples/hitl/minimal/minimal.py | 3 +- examples/hitl/pick_throw_vr/pick_throw_vr.py | 35 ++++----- examples/hitl/rearrange/rearrange.py | 23 ++++-- .../habitat_hitl/_internal/hitl_driver.py | 40 +++++----- .../habitat_hitl/app_states/app_service.py | 76 +++++++++++-------- 6 files changed, 105 insertions(+), 75 deletions(-) diff --git a/examples/hitl/basic_viewer/basic_viewer.py b/examples/hitl/basic_viewer/basic_viewer.py index d911734c0d..5076edc88c 100644 --- a/examples/hitl/basic_viewer/basic_viewer.py +++ b/examples/hitl/basic_viewer/basic_viewer.py @@ -7,6 +7,7 @@ import hydra import magnum as mn +from habitat_hitl.app_states.app_service import AppService from habitat_hitl.app_states.app_state_abc import AppState from habitat_hitl.core.gui_input import GuiInput from habitat_hitl.core.hitl_main import hitl_main @@ -18,7 +19,7 @@ class AppStateBasicViewer(AppState): def __init__( self, - app_service, + app_service: AppService, ): self._app_service = app_service self._gui_input = self._app_service.gui_input diff --git a/examples/hitl/minimal/minimal.py b/examples/hitl/minimal/minimal.py index f9e5ce8d5b..f5530d65da 100644 --- a/examples/hitl/minimal/minimal.py +++ b/examples/hitl/minimal/minimal.py @@ -7,6 +7,7 @@ import hydra import magnum +from habitat_hitl.app_states.app_service import AppService from habitat_hitl.app_states.app_state_abc import AppState from habitat_hitl.core.gui_input import GuiInput from habitat_hitl.core.hitl_main import hitl_main @@ -19,7 +20,7 @@ class AppStateMinimal(AppState): a fixed overhead camera. """ - def __init__(self, app_service): + def __init__(self, app_service: AppService): self._app_service = app_service def sim_update(self, dt, post_sim_update_dict): diff --git a/examples/hitl/pick_throw_vr/pick_throw_vr.py b/examples/hitl/pick_throw_vr/pick_throw_vr.py index 4c819108e4..f07fca46b0 100644 --- a/examples/hitl/pick_throw_vr/pick_throw_vr.py +++ b/examples/hitl/pick_throw_vr/pick_throw_vr.py @@ -4,13 +4,14 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import Final +from typing import Any, Final, List, Optional import hydra import magnum as mn import numpy as np from habitat.datasets.rearrange.navmesh_utils import get_largest_island_index +from habitat_hitl.app_states.app_service import AppService from habitat_hitl.app_states.app_state_abc import AppState from habitat_hitl.core.gui_input import GuiInput from habitat_hitl.core.hitl_main import hitl_main @@ -49,24 +50,24 @@ class AppStatePickThrowVr(AppState): See VR_HITL.md for instructions on controlling the human from a VR device. """ - def __init__(self, app_service): + def __init__(self, app_service: AppService): self._app_service = app_service - self._gui_agent_ctrl = self._app_service.gui_agent_controller + self._gui_agent_ctrl: Any = self._app_service.gui_agent_controller self._can_grasp_place_threshold = ( self._app_service.hitl_config.can_grasp_place_threshold ) - self._cam_transform = None - self._held_target_obj_idx = None - self._recent_reach_pos = None - self._paused = False - self._hide_gui_text = False + self._cam_transform: Optional[mn.Matrix4] = None + self._held_target_obj_idx: Optional[int] = None + self._recent_reach_pos: Optional[mn.Vector3] = None + self._paused: bool = False + self._hide_gui_text: bool = False # Index of the remote-controlled hand holding an object - self._remote_held_hand_idx = None + self._remote_held_hand_idx: Optional[int] = None # will be set in on_environment_reset - self._target_obj_ids = None + self._target_obj_ids: Optional[List[str]] = None self._camera_helper = CameraHelper( self._app_service.hitl_config, self._app_service.gui_input @@ -92,9 +93,9 @@ def __init__(self, app_service): self._gui_agent_ctrl.line_renderer = app_service.line_render - self._is_remote_active_toggle = False - self.count_tsteps_stop = 0 - self._has_grasp_preview = False + self._is_remote_active_toggle: bool = False + self._count_tsteps_stop: int = 0 + self._has_grasp_preview: bool = False def _is_remote_active(self): return self._is_remote_active_toggle @@ -132,7 +133,7 @@ def on_environment_reset(self, episode_recorder_dict): ) self._camera_helper.update(self._get_camera_lookat_pos(), dt=0) - self.count_tsteps_stop = 0 + self._count_tsteps_stop = 0 human_pos = ( self.get_sim() @@ -300,16 +301,16 @@ def _update_grasping_and_set_act_hints_remote(self): # Count number of steps since we stopped, this is to reduce jitter # with IK if distance_multiplier == 0: - self.count_tsteps_stop += 1 + self._count_tsteps_stop += 1 else: - self.count_tsteps_stop = 0 + self._count_tsteps_stop = 0 reach_pos = None hand_idx = None if ( self._held_target_obj_idx is not None and distance_multiplier == 0.0 - and self.count_tsteps_stop > MIN_STEPS_STOP + and self._count_tsteps_stop > MIN_STEPS_STOP ): reach_pos = self._get_target_object_position( self._held_target_obj_idx diff --git a/examples/hitl/rearrange/rearrange.py b/examples/hitl/rearrange/rearrange.py index eda0c6f14b..0f22b78c05 100644 --- a/examples/hitl/rearrange/rearrange.py +++ b/examples/hitl/rearrange/rearrange.py @@ -4,10 +4,13 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +from typing import Any, List, Optional + import hydra import magnum as mn import numpy as np +from habitat_hitl.app_states.app_service import AppService from habitat_hitl.app_states.app_state_abc import AppState from habitat_hitl.app_states.app_state_tutorial import AppStateTutorial from habitat_hitl.core.gui_input import GuiInput @@ -40,10 +43,10 @@ class AppStateRearrange(AppState): def __init__( self, - app_service, + app_service: AppService, ): self._app_service = app_service - self._gui_agent_ctrl = self._app_service.gui_agent_controller + self._gui_agent_ctrl: Any = self._app_service.gui_agent_controller # cache items from config; config is expensive to access at runtime config = self._app_service.config @@ -55,15 +58,19 @@ def __init__( self._app_service.hitl_config.can_grasp_place_threshold ) - self._cam_transform = None + self._cam_transform: Optional[mn.Matrix4] = None - self._held_target_obj_idx = None - self._num_remaining_objects = None # resting, not at goal location yet - self._num_busy_objects = None # currently held by non-gui agents + self._held_target_obj_idx: Optional[int] = None + + # resting, not at goal location yet + self._num_remaining_objects: Optional[int] = None + + # currently held by non-gui agents + self._num_busy_objects: Optional[int] = None # will be set in on_environment_reset - self._target_obj_ids = None - self._goal_positions = None + self._target_obj_ids: Optional[List[str]] = None + self._goal_positions: Optional[List[mn.Vector3]] = None self._camera_helper = CameraHelper( self._app_service.hitl_config, self._app_service.gui_input diff --git a/habitat-hitl/habitat_hitl/_internal/hitl_driver.py b/habitat-hitl/habitat_hitl/_internal/hitl_driver.py index e842cf16d9..d66f761d76 100644 --- a/habitat-hitl/habitat_hitl/_internal/hitl_driver.py +++ b/habitat-hitl/habitat_hitl/_internal/hitl_driver.py @@ -175,26 +175,30 @@ def local_end_episode(do_reset=False): if self.network_server_enabled: self._client_message_manager = ClientMessageManager() - self._app_service = AppService( - config, - self._hitl_config, - gui_input, - self._remote_gui_input, - line_render, - text_drawer, - lambda: self._viz_anim_fraction, - self.habitat_env, - self.get_sim(), - lambda: self._compute_action_and_step_env(), - self._step_recorder, - lambda: self._get_recent_metrics(), - local_end_episode, - lambda: self._set_cursor_style, - self._episode_helper, - self._client_message_manager, + gui_agent_controller: Any = ( self.ctrl_helper.get_gui_agent_controller() if self.ctrl_helper - else None, + else None + ) + + self._app_service = AppService( + config=config, + hitl_config=self._hitl_config, + gui_input=gui_input, + remote_gui_input=self._remote_gui_input, + line_render=line_render, + text_drawer=text_drawer, + get_anim_fraction=lambda: self._viz_anim_fraction, + env=self.habitat_env, + sim=self.get_sim(), + compute_action_and_step_env=lambda: self._compute_action_and_step_env(), + step_recorder=self._step_recorder, + get_metrics=lambda: self._get_recent_metrics(), + end_episode=local_end_episode, + set_cursor_style=self._set_cursor_style, + episode_helper=self._episode_helper, + client_message_manager=self._client_message_manager, + gui_agent_controller=gui_agent_controller, ) self._app_state: AppState = None diff --git a/habitat-hitl/habitat_hitl/app_states/app_service.py b/habitat-hitl/habitat_hitl/app_states/app_service.py index 229a2225a2..e06489fc20 100644 --- a/habitat-hitl/habitat_hitl/app_states/app_service.py +++ b/habitat-hitl/habitat_hitl/app_states/app_service.py @@ -4,28 +4,44 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. +from __future__ import annotations + +from typing import Callable, Optional + +from habitat import Env +from habitat.tasks.rearrange.rearrange_sim import RearrangeSim +from habitat_hitl.core.client_message_manager import ClientMessageManager +from habitat_hitl.core.gui_input import GuiInput +from habitat_hitl.core.remote_gui_input import RemoteGuiInput +from habitat_hitl.core.serialize_utils import BaseRecorder +from habitat_hitl.core.text_drawer import AbstractTextDrawer +from habitat_hitl.environment.controllers.controller_abc import GuiController +from habitat_hitl.environment.episode_helper import EpisodeHelper +from habitat_sim.gfx import DebugLineRender + # Helpers to provide to AppState classes, provided by the underlying SandboxDriver class AppService: def __init__( self, + *, config, hitl_config, - gui_input, - remote_gui_input, - line_render, - text_drawer, - get_anim_fraction, - env, - sim, - compute_action_and_step_env, - step_recorder, - get_metrics, - end_episode, - set_cursor_style, - episode_helper, - client_message_manager, - gui_agent_controller, + gui_input: GuiInput, + remote_gui_input: RemoteGuiInput, + line_render: DebugLineRender, + text_drawer: AbstractTextDrawer, + get_anim_fraction: Callable, + env: Env, + sim: RearrangeSim, + compute_action_and_step_env: Callable, + step_recorder: BaseRecorder, + get_metrics: Callable, + end_episode: Callable, + set_cursor_style: Callable, + episode_helper: EpisodeHelper, + client_message_manager: ClientMessageManager, + gui_agent_controller: Optional[GuiController], ): self._config = config self._hitl_config = hitl_config @@ -54,61 +70,61 @@ def hitl_config(self): return self._hitl_config @property - def gui_input(self): + def gui_input(self) -> GuiInput: return self._gui_input @property - def remote_gui_input(self): + def remote_gui_input(self) -> RemoteGuiInput: return self._remote_gui_input @property - def line_render(self): + def line_render(self) -> DebugLineRender: return self._line_render @property - def text_drawer(self): + def text_drawer(self) -> AbstractTextDrawer: return self._text_drawer @property - def get_anim_fraction(self): + def get_anim_fraction(self) -> Callable: return self._get_anim_fraction @property - def env(self): + def env(self) -> Env: return self._env @property - def sim(self): + def sim(self) -> RearrangeSim: return self._sim @property - def compute_action_and_step_env(self): + def compute_action_and_step_env(self) -> Callable: return self._compute_action_and_step_env @property - def step_recorder(self): + def step_recorder(self) -> BaseRecorder: return self._step_recorder @property - def get_metrics(self): + def get_metrics(self) -> Callable: return self._get_metrics @property - def end_episode(self): + def end_episode(self) -> Callable: return self._end_episode @property - def set_cursor_style(self): + def set_cursor_style(self) -> Callable: return self._set_cursor_style @property - def episode_helper(self): + def episode_helper(self) -> EpisodeHelper: return self._episode_helper @property - def client_message_manager(self): + def client_message_manager(self) -> ClientMessageManager: return self._client_message_manager @property - def gui_agent_controller(self): + def gui_agent_controller(self) -> Optional[GuiController]: return self._gui_agent_controller