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 typing import Any
  7from collections import deque
  8
  9from overrides import overrides
 10import numpy as np
 11import raylib as rl
 12from microecs import World, EntityId, Component
 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, logger, make_arr, FixedSizeDict
 19from robolib.constants import FPV_HEIGHT, FPV_WIDTH, DEFAULT_BACKGROUND, TICK_STATS_LEN
 20from robolib.traits import Restorable, Serializable, Drawable
 21from robolib.plugins_manager import PluginsManager
 22from robolib.components import ROBOLIB_COMPONENTS, ColliderKinds, HasModel, HasPose, HasFPV, HasCollision
 23from robolib.entities import add_entity
 24
 25from plugins import REGISTERED_PLUGINS
 26
 27# simulator-specific components (e.g. not generic which whould live in robosim/components).
 28# generic ones should have the 'Has' prefix (e.g. HasPose), while non-generic ones usually don't (e.g. GroundFloor)
 29
 30class GroundFloor(Component):
 31    """component used to tag the floor so we don't render it in collision mode"""
 32
 33ALL_COMPONENTS = [*ROBOLIB_COMPONENTS, GroundFloor]
 34
 35# serialization / build entities
 36
 37def world_to_dict(world: World) -> dict[str, Any]:
 38    """Serialize the world. Goes through all the entities and their components and converts the serializables to dict"""
 39    res = {"entities": [], "components": world.component_names, "extra_metadata": world.extra_metadata}
 40    for eid in world.live_entities.keys():
 41        res["entities"].append({"entity_id": eid, **world.get_entity(eid).to_dict(serialization_field="serializable")})
 42    return res
 43
 44def world_from_dict(data: dict[str, Any]) -> World:
 45    """Creates a world from a serialized representation e.g. from world_to_dict()"""
 46    components = [{c.__name__: c for c in ALL_COMPONENTS}[c] for c in data["components"]]
 47    world = World(components=components, extra_metadata=data["extra_metadata"])
 48    for entity in data["entities"]:
 49        components = [world.component_name_to_type[c] for c in entity["components"]]
 50        add_entity(world, components, **entity["data"])
 51    return world
 52
 53# Simulator State object. All relevant state data stays here. This will be serialized and hashed for DST in the future.
 54
 55@dataclass(kw_only=True)
 56class SimState(Serializable):
 57    """The state of the scene including the cameras, robot and so on"""
 58    world_camera: Camera | None = None
 59    topdown_camera: Camera | None = None
 60    active_camera_type: str = "world" # Which camera is active based on it's type
 61    fpv_ix: int = 0
 62    wireframe_mode: bool = False
 63    collision_render_mode: bool = False
 64    collision_cell_size: Point3D = field(default_factory=lambda: make_arr(30, 30, 30))
 65    display_uav_trace: bool = False
 66    uav_traces: dict[EntityId, FixedSizeDict[tuple[float, float, float], bool]] = field(default_factory=dict)
 67    uav_traces_max_len: int = 1000
 68
 69    def __post_init__(self):
 70        if self.world_camera is None:
 71            self.world_camera = Camera(make_arr(15, 12, 15), make_arr(0, 2, 0), make_arr(0, 1, 0),
 72                                       fovy=45, projection=rl.CAMERA_PERSPECTIVE)
 73        if self.topdown_camera is None:
 74            self.topdown_camera = Camera(make_arr(0, 30, 0), make_arr(0, 0, 0), make_arr(0, 0, -1),
 75                                         fovy=20, projection=rl.CAMERA_ORTHOGRAPHIC)
 76
 77    @property
 78    def active_camera_label(self) -> str:
 79        """returns a string representation of the active camera for HUD"""
 80        match self.active_camera_type:
 81            case "world": return "WORLD CAMERA (F1)"
 82            case "topdown": return "TOP-DOWN CAMERA (F2)"
 83            case "fpv": return "FPV CAMERA (F3)"
 84            case _: raise ValueError(self.active_camera_type)
 85
 86    @overrides
 87    def to_dict(self) -> dict:
 88        return {
 89            "world_camera": self.world_camera.to_dict(),
 90            "topdown_camera": self.topdown_camera.to_dict(),
 91            "active_camera_type": self.active_camera_type,
 92            "fpv_ix": self.fpv_ix,
 93            "wireframe_mode": self.wireframe_mode,
 94            "collision_render_mode": self.collision_render_mode,
 95            "collision_cell_size": self.collision_cell_size.tolist(),
 96            "display_uav_trace": self.display_uav_trace,
 97            "uav_traces": FixedSizeDict({k: list(v.keys()) for k, v in self.uav_traces.items()},
 98                                        maxlen=self.uav_traces_max_len),
 99            "uav_traces_max_len": self.uav_traces_max_len,
100        }
101
102    @staticmethod
103    @overrides
104    def from_dict(state: dict) -> SimState:
105        return SimState(
106            world_camera=Camera.from_dict(state["world_camera"]),
107            topdown_camera=Camera.from_dict(state["topdown_camera"]),
108            active_camera_type=state["active_camera_type"],
109            fpv_ix=state["fpv_ix"],
110            wireframe_mode=state["wireframe_mode"],
111            collision_render_mode=state["collision_render_mode"],
112            collision_cell_size=np.float32(state["collision_cell_size"]),
113            display_uav_trace=state["display_uav_trace"],
114            uav_traces={int(eid): FixedSizeDict({tuple(pt): True for pt in points}, maxlen=state["uav_traces_max_len"])
115                        for eid, points in state["uav_traces"].items()},
116            uav_traces_max_len=state["uav_traces_max_len"],
117        )
118
119class Simulator(Restorable):
120    """The simulator object. Contains the scene (objects/robots), state and events. Wires to/from_dict as well."""
121    def __init__(self, world: World, state: SimState, connection_manager: ConnectionManager,
122                 plugins: list[str], protocol: Protocol):
123        self.world = world
124        self.state = state
125        self.connection_manager = connection_manager
126        self.plugins = plugins
127        self.protocol = protocol
128
129        self.plugins_manager = PluginsManager(registered_plugins=REGISTERED_PLUGINS, plugin_names=plugins)
130        self.physics_ticks_stats = deque(maxlen=TICK_STATS_LEN)
131        self.render_ticks_stats = deque(maxlen=TICK_STATS_LEN)
132
133        self.world.update()
134        self.robot_eids: list[EntityId] = self.world.query(HasFPV).entity_ids.tolist()
135        self.robot_eid_to_channel_ix: dict[EntityId, int] = {}
136        self._add_fpv_to_robots()
137
138        self._all_commands: set[str] | None = None # plugins + core commands cache
139
140    @property
141    def active_camera(self) -> Camera:
142        """The current active camera out of the 3: world, uav or topdown"""
143        match self.state.active_camera_type:
144            case "world": return self.state.world_camera
145            case "topdown": return self.state.topdown_camera
146            case "fpv": return self.world.get_entity(
147                self.world.query(HasFPV).entity_ids[self.state.fpv_ix]).fpv_camera[0]
148            case _: raise ValueError(self.state.active_camera_type)
149
150    @property
151    def all_commands(self) -> set[str]:
152        """All the commands (core and plugins) available in the simulator. Must be updated on plugins/world changes"""
153        if self._all_commands is None:
154            self._all_commands = set(self.protocol.endpoints) | self.plugins_manager.all_endpoints
155        return self._all_commands
156
157    def render_main_camera(self, extra_drawables: list[Drawable] | None):
158        """renders the main camera (the raylib UI) of the simulator"""
159        rl.BeginDrawing()
160        rl.ClearBackground(DEFAULT_BACKGROUND)
161
162        rl.BeginMode3D(self.active_camera.camera[0])
163
164        if self.state.wireframe_mode:
165            rl.rlEnableWireMode()
166
167        if self.state.collision_render_mode:
168            self._render_collision_mode(extra_drawables)
169        else:
170            entity_id = self.robot_eids[self.state.fpv_ix] if self.state.active_camera_type == "fpv" else None
171            self._render(entity_id=entity_id, extra_drawables=extra_drawables)
172
173        if self.state.display_uav_trace:
174            for traces in self.state.uav_traces.values():
175                for trace in traces.keys():
176                    rl.DrawSphere(trace, 0.05, rl.GREEN)
177
178        rl.EndMode3D()
179
180        if self.state.wireframe_mode:
181            rl.rlDisableWireMode()
182
183        n_robots = len(self.world.query(HasFPV))
184        msg = f"{self.state.active_camera_label}\nObjects: {len(self.world)} (robots: {n_robots})"
185        rl.DrawText(msg.encode(), 10, 10, 20, rl.DARKGRAY)
186        rl.DrawFPS(rl.GetScreenWidth() - 80, 10)
187        rl.EndDrawing()
188
189    def render_fpv_cameras(self, extra_drawables: list[Drawable] | None):
190        """Renders each robot's FPV camera but only if it's streaming. Calls self.draw() for scene content."""
191        for robot_eid in self.robot_eids:
192            if robot_eid not in self.robot_eid_to_channel_ix:
193                continue
194            entity = self.world.get_entity(robot_eid)
195            fpv_texture: "rl.RenderTexture" = entity.fpv_texture.item()
196            fpv_data: FPVData = entity.fpv_data.item()
197            camera: Camera = entity.fpv_camera.item()
198
199            rl.BeginTextureMode(fpv_texture)
200            rl.ClearBackground(DEFAULT_BACKGROUND)
201            rl.BeginMode3D(camera.camera[0])
202            self._render(entity_id=robot_eid, extra_drawables=extra_drawables)
203            rl.EndMode3D()
204            rl.EndTextureMode()
205
206            fpv_img = rl.LoadImageFromTexture(fpv_texture.texture)
207            fpv_img_ptr = rl.ffi.new("Image *", fpv_img)
208            rl.ImageFlipVertical(fpv_img_ptr)
209            with fpv_data.lock:
210                fpv_data.frame = bytes(rl.ffi.buffer(fpv_img_ptr.data, len(fpv_data.frame)))
211                fpv_data.frame_id += 1
212                fpv_data._frame_compressed = None # pylint: disable=protected-access
213            rl.UnloadImage(fpv_img_ptr[0])
214
215    def assign_channel_to_first_free_robot(self, channel_idx: int) -> int | None:
216        """Gets the first free robot eid. Called from protocol._handle_message (slow). TODO: Use a lock ?"""
217        for robot_ix, robot_eid in enumerate(self.robot_eids):
218            if robot_eid not in self.robot_eid_to_channel_ix:
219                self.robot_eid_to_channel_ix[robot_eid] = channel_idx
220                return robot_ix
221        return None
222
223    def client_on_disconnect(self, client: RobosimClient):
224        """Clears the robot that was assigned to this connected client. TODO: use a lock ?"""
225        if client.state == RobosimClientState.CONNECTED:
226            robot_ix = self.robot_eids.index(client.robot_eid)
227            logger.debug(f"Releasing robot {robot_ix} (eid: {client.robot_eid}, channel: {client.channel.idx})")
228            self.robot_eid_to_channel_ix.pop(client.robot_eid)
229            client.state = RobosimClientState.STAGED
230
231    @overrides
232    def to_dict(self) -> dict:
233        return {
234            "world": world_to_dict(self.world),
235            "state": self.state.to_dict(),
236            "plugins": self.plugins_manager.to_dict(),
237        }
238
239    @overrides
240    def load_state_dict(self, state: dict):
241        state = deepcopy(state) # this is mostly for INIT_STATE as the code below mutates it while loading
242
243        self._load_world_from_dict(state["world"])
244        self.state = SimState.from_dict(state["state"])
245
246        old_names = set(self.plugins)
247        for dropped in old_names - state["plugins"].keys():
248            logger.warning(f"Plugin '{dropped}' was live but absent from loaded state. Dropping.")
249        self.plugins = list(state["plugins"])
250        self.plugins_manager = PluginsManager.from_dict(REGISTERED_PLUGINS, state["plugins"])
251        self._all_commands = None # Reset this so all_commands is created again and we don't re-use the cache.
252
253    def _add_fpv_to_robots(self):
254        for eid in self.robot_eids: # FPV data is not part of ECS
255            camera = Camera(make_arr(0, 0, 0), make_arr(0, 0, 0), make_arr(0, 0, 0), 60.0, rl.CAMERA_PERSPECTIVE)
256            fpv_data = FPVData(frame=np.zeros((FPV_HEIGHT * FPV_WIDTH * 4,), "uint8"),
257                               frame_shape=(FPV_HEIGHT, FPV_WIDTH, 4))
258            robot = self.world.get_entity(eid)
259            # TODO(microecs-30): set_data(fpv_camera=fpv_camera) and auto-convert on microecs side.
260            robot.set_data(fpv_camera=np.array([camera], "object"))
261            robot.set_data(fpv_texture=np.array([rl.LoadRenderTexture(FPV_WIDTH, FPV_HEIGHT)], "object")) # TODO: dealoc
262            robot.set_data(fpv_data=np.array([fpv_data], "object"))
263        self.world.update()
264
265    def _load_world_from_dict(self, world_state: dict):
266        self.world = world_from_dict(world_state)
267        self.world.update()
268
269        self.robot_eids = self.world.query(HasFPV).entity_ids.tolist()
270        self._add_fpv_to_robots()
271
272        logger.info(f"Loaded world from state: {self.world}")
273
274    def _render(self, entity_id: int | None, extra_drawables: list[Drawable] | None): #noqa
275        """Renders the scene and all its objects + robots. Called from within a rl.BeginMode3D() block."""
276        qr = self.world.query(HasModel, HasPose)
277        for i, (model, pose, scale) in enumerate(zip(qr.model, qr.pose, qr.scale)):
278            if entity_id is not None and entity_id == qr.entity_ids[i]: # if in FPV mode then don't draw yourself
279                continue
280            rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
281
282        for drawable in (extra_drawables or []): # e.g. plugins
283            drawable.draw(self.world, entity_id)
284
285    def _render_collision_mode(self, extra_drawables: list[Drawable] | None):
286        """render the scene in 'collision mode' which means we don't render the meshes of collidables, but the shapes"""
287        qr = self.world.query(HasModel, HasPose, exclude=[HasCollision])
288        for model, pose, scale in zip(qr.model, qr.pose, qr.scale):
289            rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
290
291        qr = self.world.query(HasModel, HasPose, HasCollision, exclude=[GroundFloor])
292        collider_size = qr.collider_bbox[:, 1] - qr.collider_bbox[:, 0]
293        for pose, kind, radius, size, colliding in zip(qr.pose, qr.collider_kind, qr.collider_radii,
294                                                       collider_size, qr.is_colliding):
295            color = rl.RED if colliding else rl.GREEN
296            if kind == ColliderKinds.SPHERE:
297                rl.DrawSphere(pose[0:3, 3].tolist(), radius.item(), color)
298            elif kind == ColliderKinds.AABB:
299                rl.DrawCube(pose[0:3, 3].tolist(), *size.tolist(), color)
300            else:
301                raise NotImplementedError(kind)
302
303        qr = self.world.query(HasModel, HasPose, GroundFloor)
304        for model, pose, scale in zip(qr.model, qr.pose, qr.scale):
305            rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
306
307        for drawable in (extra_drawables or []): # e.g. plugins
308            drawable.draw(self.world, None)
309
310    def __repr__(self):
311        return f"[Simulator]\n{self.world}\n{self.plugins_manager}\n{self.connection_manager}\n{self.protocol}"