refactor(simulator): extract habitat simulator and visualization modules

mini-nav/simulator/views.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any, Callable, Iterable, Sequence
+
+from rich.progress import track
+
+RoomViewsByRoom = dict[str, list[Any]]
+ProgressTrack = Callable[[Iterable[Any], str], Iterable[Any]]
+
+
+def collect_room_views_by_room(
+    agent: Any,
+    sim: Any,
+    room_nodes: Sequence[Any],
+    views_per_room: int,
+    *,
+    habitat_sim_module: Any | None = None,
+    sensor_uuid: str = "color_sensor",
+    turn_action: str = "turn_left",
+    progress_description: str = "Collecting room views",
+    progress_track: ProgressTrack = track,
+) -> RoomViewsByRoom:
+    if views_per_room <= 0:
+        raise ValueError("views_per_room must be greater than 0")
+
+    if habitat_sim_module is None:
+        habitat_sim_module = import_module("habitat_sim")
+
+    all_room_views: RoomViewsByRoom = {}
+    for room_node in progress_track(room_nodes, progress_description):
+        agent_state = habitat_sim_module.AgentState()
+        agent_state.position = room_node.center.copy()
+        agent.set_state(agent_state)
+
+        room_views = []
+        for _ in range(views_per_room):
+            observations = sim.get_sensor_observations()
+            room_views.append(observations[sensor_uuid])
+            sim.step(turn_action)
+
+        all_room_views[room_node.room_id] = room_views
+
+    return all_room_views