Files
Mini-Nav/mini-nav/simulator/habitat.py

74 lines
2.4 KiB
Python

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 = 512
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()