refactor(simulator): extract habitat simulator and visualization modules

mini-nav/simulator/habitat.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from importlib import import_module
+from typing import Any
+
+
+@dataclass(frozen=True)
+class HabitatSimulatorConfig:
+    scene_path: str
+    views_per_room: int = 6
+    image_size: int = 256
+    sensor_height: float = 1.5
+    move_forward_step: float = 0.25
+    enable_physics: bool = False
+    sensor_uuid: str = "color_sensor"
+    agent_id: int = 0
+
+
+def create_habitat_simulator(
+    config: HabitatSimulatorConfig,
+    habitat_sim_module: Any | None = None,
+) -> tuple[Any, Any]:
+    if config.views_per_room <= 0:
+        raise ValueError("views_per_room must be greater than 0")
+
+    if config.image_size <= 0:
+        raise ValueError("image_size must be greater than 0")
+
+    if config.move_forward_step <= 0:
+        raise ValueError("move_forward_step must be greater than 0")
+
+    if habitat_sim_module is None:
+        habitat_sim_module = import_module("habitat_sim")
+
+    sim_cfg = habitat_sim_module.SimulatorConfiguration()
+    sim_cfg.scene_id = config.scene_path
+    sim_cfg.enable_physics = config.enable_physics
+
+    agent_cfg = habitat_sim_module.agent.AgentConfiguration()
+    rgb_sensor_spec = habitat_sim_module.CameraSensorSpec()
+    rgb_sensor_spec.uuid = config.sensor_uuid
+    rgb_sensor_spec.sensor_type = habitat_sim_module.SensorType.COLOR
+    rgb_sensor_spec.resolution = [config.image_size, config.image_size]
+    rgb_sensor_spec.position = [0.0, config.sensor_height, 0.0]
+    agent_cfg.sensor_specifications = [rgb_sensor_spec]
+
+    turn_angle = 360.0 / config.views_per_room
+    agent_cfg.action_space = {
+        "move_forward": habitat_sim_module.agent.ActionSpec(
+            "move_forward",
+            habitat_sim_module.agent.ActuationSpec(amount=config.move_forward_step),
+        ),
+        "turn_left": habitat_sim_module.agent.ActionSpec(
+            "turn_left",
+            habitat_sim_module.agent.ActuationSpec(amount=turn_angle),
+        ),
+        "turn_right": habitat_sim_module.agent.ActionSpec(
+            "turn_right",
+            habitat_sim_module.agent.ActuationSpec(amount=turn_angle),
+        ),
+    }
+
+    simulator_cfg = habitat_sim_module.Configuration(sim_cfg, [agent_cfg])
+    simulator = habitat_sim_module.Simulator(simulator_cfg)
+    agent = simulator.initialize_agent(config.agent_id)
+    return simulator, agent
+
+
+def close_habitat_simulator(simulator: Any) -> None:
+    close = getattr(simulator, "close", None)
+    if callable(close):
+        close()