robosim.simulator_singleton

  1"""simulator_singleton.py - The simulator singleton that holds the world, the sim state, plugins, protocol etc."""
  2
  3from __future__ import annotations
  4from dataclasses import dataclass, field
  5from copy import deepcopy
  6from collections import deque
  7from enum import StrEnum
  8
  9from overrides import overrides
 10import numpy as np
 11import raylib as rl
 12from microecs import World, EntityId
 13from microspec import Protocol
 14from micronetcode import ConnectionManager
 15
 16from robolib.camera import Camera
 17from robolib.netcode import RobosimClient, RobosimClientState
 18from robolib.utils import Point3D,FPVData, FixedSizeDict, Clock, logger, make_arr
 19from robolib.constants import FPV_HEIGHT, FPV_WIDTH, TICK_STATS_LEN, MAX_SUBTICKS_PER_RENDER_TICK, DT
 20from robolib.traits import Restorable, Serializable
 21from robolib.plugins_manager import PluginsManager
 22from robolib.components import ROBOLIB_COMPONENTS, HasModel, HasFPV, HasCollision
 23from robolib.physics import make_collision_cell_size
 24from robolib.serialization import world_from_dict, world_to_dict
 25
 26from plugins import REGISTERED_PLUGINS
 27
 28# simulator-specific components (e.g. not generic which whould live in robosim/components).
 29# generic ones should have the 'Has' prefix (e.g. HasPose), while non-generic ones usually don't (e.g. GroundFloor)
 30
 31ALL_COMPONENTS = [*ROBOLIB_COMPONENTS]
 32
 33# Simulator State object. All relevant state data stays here. This will be serialized and hashed for DST in the future.
 34
 35class CameraTypes(StrEnum):
 36    """the types of cameras"""
 37    WORLD  = "world"
 38    TOPDOW = "topdown"
 39    FPV    = "fpv"
 40
 41@dataclass(kw_only=True)
 42class SimState(Serializable):
 43    """The state of the scene including the cameras, robot and so on"""
 44    world_camera: Camera | None = None
 45    topdown_camera: Camera | None = None
 46    active_camera_type: CameraTypes = CameraTypes.WORLD
 47    fpv_ix: int = 0
 48    wireframe_mode: bool = False
 49    collision_render_mode: bool = False
 50    display_uav_trace: bool = False
 51    uav_traces: dict[EntityId, FixedSizeDict[tuple[float, float, float], bool]] = field(default_factory=dict)
 52    uav_traces_max_len: int = 1000
 53    must_exit: bool = False
 54
 55    def __post_init__(self):
 56        if self.world_camera is None:
 57            self.world_camera = Camera(ptu={"position": make_arr(15, 12, 15), "target": make_arr(0, 2, 0),
 58                                        "up": make_arr(0, 1, 0)}, fovy=45, projection=rl.CAMERA_PERSPECTIVE,
 59                                        label="World Camera")
 60        if self.topdown_camera is None:
 61            self.topdown_camera = Camera(ptu={"position": make_arr(0, 30, 0), "target": make_arr(0, 0, 0),
 62                                          "up": make_arr(0, 0, -1)}, fovy=20, projection=rl.CAMERA_ORTHOGRAPHIC,
 63                                          label="Topdown Camera")
 64
 65    @overrides
 66    def to_dict(self) -> dict:
 67        return {
 68            "world_camera": self.world_camera.to_dict(),
 69            "topdown_camera": self.topdown_camera.to_dict(),
 70            "active_camera_type": self.active_camera_type,
 71            "fpv_ix": self.fpv_ix,
 72            "wireframe_mode": self.wireframe_mode,
 73            "collision_render_mode": self.collision_render_mode,
 74            "display_uav_trace": self.display_uav_trace,
 75            "uav_traces": FixedSizeDict({k: list(v.keys()) for k, v in self.uav_traces.items()},
 76                                        maxlen=self.uav_traces_max_len),
 77            "uav_traces_max_len": self.uav_traces_max_len,
 78            "must_exit": self.must_exit,
 79        }
 80
 81    @staticmethod
 82    @overrides
 83    def from_dict(state: dict) -> SimState:
 84        return SimState(
 85            world_camera=Camera.from_dict(state["world_camera"]),
 86            topdown_camera=Camera.from_dict(state["topdown_camera"]),
 87            active_camera_type=state["active_camera_type"],
 88            fpv_ix=state["fpv_ix"],
 89            wireframe_mode=state["wireframe_mode"],
 90            collision_render_mode=state["collision_render_mode"],
 91            display_uav_trace=state["display_uav_trace"],
 92            uav_traces={int(eid): FixedSizeDict({tuple(pt): True for pt in points}, maxlen=state["uav_traces_max_len"])
 93                        for eid, points in state["uav_traces"].items()},
 94            uav_traces_max_len=state["uav_traces_max_len"],
 95            must_exit=state["must_exit"],
 96        )
 97
 98class Simulator(Restorable):
 99    """The simulator object. Contains the scene (objects/robots), state and events. Wires to/from_dict as well."""
100    def __init__(self, world: World, state: SimState, connection_manager: ConnectionManager,
101                 plugins: list[str], protocol: Protocol):
102        self.world = world
103        self.state = state
104        self.connection_manager = connection_manager
105        self.plugins = plugins
106        self.protocol = protocol
107
108        self.clock = Clock(dt=DT, max_ticks=MAX_SUBTICKS_PER_RENDER_TICK)
109        self.plugins_manager = PluginsManager(registered_plugins=REGISTERED_PLUGINS, plugin_names=plugins)
110        self.physics_ticks_stats = deque(maxlen=TICK_STATS_LEN)
111        self.render_ticks_stats = deque(maxlen=TICK_STATS_LEN)
112
113        self.world.update()
114        self.robot_eids: list[EntityId] = self.world.query(HasFPV).entity_ids.tolist()
115        self.robot_eid_to_channel_ix: dict[EntityId, int] = {}
116        self._add_fpv_to_robots()
117
118        self._all_commands: set[str] | None = None # plugins + core commands cache
119        self._collision_cell_size: Point3D | None = None
120
121    @property
122    def collision_cell_size(self) -> Point3D:
123        """The collision cell size (for the hash grid). If not set, computes the cell size based on the world"""
124        if self._collision_cell_size is None:
125            self.world.update() # update so we use all entities (even those recently added via add_entity)
126            self._collision_cell_size = make_collision_cell_size(entities=self.world.query(HasModel, HasCollision))
127        return self._collision_cell_size
128
129    @collision_cell_size.setter
130    def collision_cell_size(self, size: Point3D):
131        self._collision_cell_size = size
132
133    @property
134    def active_camera(self) -> Camera:
135        """The current active camera out of the 3: world, uav or topdown"""
136        match self.state.active_camera_type:
137            case "world": return self.state.world_camera
138            case "topdown": return self.state.topdown_camera
139            case "fpv": return self.world.get_entity(
140                self.world.query(HasFPV).entity_ids[self.state.fpv_ix]).fpv_camera[0]
141            case _: raise ValueError(self.state.active_camera_type)
142
143    @property
144    def all_commands(self) -> set[str]:
145        """All the commands (core and plugins) available in the simulator. Must be updated on plugins/world changes"""
146        if self._all_commands is None:
147            self._all_commands = set(self.protocol.endpoints) | self.plugins_manager.all_endpoints
148        return self._all_commands
149
150    def assign_channel_to_first_free_robot(self, channel_idx: int) -> int | None:
151        """Gets the first free robot eid. Called from protocol._handle_message (slow). TODO: Use a lock ?"""
152        for robot_ix, robot_eid in enumerate(self.robot_eids):
153            if robot_eid not in self.robot_eid_to_channel_ix:
154                self.robot_eid_to_channel_ix[robot_eid] = channel_idx
155                return robot_ix
156        return None
157
158    def client_on_disconnect(self, client: RobosimClient):
159        """Clears the robot that was assigned to this connected client. TODO: use a lock ?"""
160        if client.state == RobosimClientState.CONNECTED:
161            robot_ix = self.robot_eids.index(client.robot_eid)
162            logger.debug(f"Releasing robot {robot_ix} (eid: {client.robot_eid}, channel: {client.channel.idx})")
163            self.robot_eid_to_channel_ix.pop(client.robot_eid)
164            client.state = RobosimClientState.STAGED
165
166    @overrides
167    def to_dict(self) -> dict:
168        return {
169            "world": world_to_dict(self.world),
170            "state": self.state.to_dict(),
171            "plugins": self.plugins_manager.to_dict(),
172        }
173
174    @overrides
175    def load_state_dict(self, state: dict):
176        state = deepcopy(state) # this is mostly for INIT_STATE as the code below mutates it while loading
177        self._collision_cell_size = None # Reset this as the world may've changed and we need new collision cell sizes
178        self._all_commands = None # Reset this so all_commands is created again and we don't re-use the cache.
179
180        self._load_world_from_dict(state["world"])
181        if "state" in state:
182            self.state = SimState.from_dict(state["state"])
183        else:
184            logger.warning("'state' key not found in the state dict. Skipping SimulatorState update.")
185
186        if "plugins" in state:
187            old_names = set(self.plugins)
188            for dropped in old_names - state["plugins"].keys():
189                logger.warning(f"Plugin '{dropped}' was live but absent from loaded state. Dropping.")
190            self.plugins = list(state["plugins"])
191            self.plugins_manager = PluginsManager.from_dict(REGISTERED_PLUGINS, state["plugins"])
192        else:
193            logger.warning("'plugins' key not found in the state dict. Skipping PluginsManager update.")
194
195    def _add_fpv_to_robots(self):
196        for eid in self.robot_eids: # FPV data is not part of ECS
197            camera = Camera(fovy=60.0, projection=rl.CAMERA_PERSPECTIVE, label="FPV")
198            fpv_data = FPVData(frame=np.zeros((FPV_HEIGHT * FPV_WIDTH * 4,), "uint8"),
199                               frame_shape=(FPV_HEIGHT, FPV_WIDTH, 4))
200            robot = self.world.get_entity(eid)
201            # TODO(microecs-30): set_data(fpv_camera=fpv_camera) and auto-convert on microecs side.
202            robot.set_data(fpv_camera=np.array([camera], "object"))
203            robot.set_data(fpv_texture=np.array([rl.LoadRenderTexture(FPV_WIDTH, FPV_HEIGHT)], "object")) # TODO: dealoc
204            robot.set_data(fpv_data=np.array([fpv_data], "object"))
205        self.world.update()
206
207    def _load_world_from_dict(self, world_state: dict):
208        self.world = world_from_dict(world_state)
209        self.world.update()
210
211        self.robot_eids = self.world.query(HasFPV).entity_ids.tolist()
212        self._add_fpv_to_robots()
213
214        logger.info(f"Loaded world from state: {self.world}")
215
216    def __repr__(self):
217        return f"[Simulator]\n{self.world}\n{self.plugins_manager}\n{self.connection_manager}\n{self.protocol}"