server.server

  1#!/usr/bin/env python3
  2"""
  3Minimal raylib-based UAV Trajectory Simulation with TCP connection to control the drone.
  4Constructs a simple obstacle world (cubes on a grid).
  5
  6Keybinds:
  7- Controls on cameras (world, topdown, FPV):
  8  - W, A, S, D (both cameras) + mouse for rotation (just for world camera).
  9  - Q, E, arrows, PgUp, PgDown (just for FPV)
 10- Misc:
 11  - F1/F2/F3: Switch to world camera (3D overview) / top-dopwn camera (2D view) / FPV camera (UAV)
 12  - ESC: Quit
 13"""
 14from __future__ import annotations
 15from dataclasses import dataclass, field
 16from pathlib import Path
 17from copy import deepcopy
 18from queue import Empty
 19from typing import Any
 20from collections import deque
 21from functools import partial
 22from argparse import ArgumentParser, Namespace
 23import traceback
 24import json
 25import sys
 26import time
 27
 28from overrides import overrides
 29import numpy as np
 30import raylib as rl
 31from microecs import World, EntityId, Component
 32from microspec import Protocol
 33
 34from robosim.camera import Camera
 35from robosim.network import ConnectionManager, ObjToSimChannel
 36from robosim.utils import (Point3D,  Clock, FPVData, logger, get_project_root, get_closest_square,
 37                           pose_from_position_target_up, pose_from_trans_euler as pte, make_arr, FixedSizeDict,
 38                           rl_get_device_and_renderer)
 39from robosim.constants import (HZ, SCREEN_HEIGHT, SCREEN_WIDTH, DT, CONNECT_CMD,
 40                               MAX_SUBTICKS_PER_RENDER_TICK, SOCKET_TIMEOUT_S, DEFAULT_BACKGROUND, TICK_STATS_LEN,
 41                               LVL2_DRAG_COEFF, LVL2_MAX_ACCELERATIONS)
 42from robosim.traits import Restorable, Serializable, Drawable
 43from robosim.plugins_manager import PluginsManager
 44from robosim.physics import make_collision_cell_size
 45from robosim.entities import make_cube, make_mesh, make_robot
 46from robosim.components import (ROBOSIM_COMPONENTS, ColliderKinds, HasModel, HasPose, HasFPV, HasChannel,
 47                                HasMotionInput, HasPhysicsLevel1, HasPhysicsLevel2, HasCollision)
 48from robosim.systems import FPVCameraSystem, PhysicsSystem
 49
 50np.set_printoptions(precision=3, linewidth=120)
 51rl.SetTraceLogLevel(rl.LOG_WARNING)
 52
 53sys.path.append(str(get_project_root() / "server/plugins"))
 54from manual_move_plugin import ManualMovePlugin # pylint: disable=wrong-import-order, import-error, wrong-import-position
 55from race_plugin import RacePlugin # pylint: disable=wrong-import-order, import-error, wrong-import-position
 56from trajectory_mission_plugin import TrajectoryMissionPlugin # pylint: disable=wrong-import-order, import-error, wrong-import-position
 57
 58# global settings and types
 59
 60RESOURCES_PATH = get_project_root() / "resources"
 61INIT_CONFIG: dict | None = None
 62UAV_TRACES_MAX_LEN = 1000
 63
 64PluginsManager.REGISTERED_PLUGINS = {  # plugin registration
 65    "manual_move": ManualMovePlugin,
 66    "trajectory_mission": TrajectoryMissionPlugin,
 67    "race": RacePlugin,
 68}
 69
 70# simulator-specific components (e.g. not generic which whould live in robosim/components).
 71# generic ones should have the 'Has' prefix (e.g. HasPose), while non-generic ones usually don't (e.g. GroundFloor)
 72
 73class GroundFloor(Component):
 74    """component used to tag the floor so we don't render it in collision mode"""
 75
 76ALL_COMPONENTS = [*ROBOSIM_COMPONENTS, GroundFloor]
 77
 78# serialization / build entities
 79
 80def world_to_dict(world: World) -> dict[str, Any]:
 81    """Serialize the world. Goes through all the entities and their components and converts the serializables to dict"""
 82    res = {"entities": [], "components": world.component_names, "extra_metadata": world.extra_metadata}
 83    for eid in world.live_entities.keys():
 84        res["entities"].append({"entity_id": eid, **world.get_entity(eid).to_dict(serialization_field="serializable")})
 85    return res
 86
 87def world_from_dict(data: dict[str, Any]) -> World:
 88    """Creates a world from a serialized representation e.g. from world_to_dict()"""
 89    components = [{c.__name__: c for c in ALL_COMPONENTS}[c] for c in data["components"]]
 90    world = World(components=components, extra_metadata=data["extra_metadata"])
 91
 92    for entity in data["entities"]:
 93        entity_data = entity["data"]
 94        entity_components = [{c.__name__: c for c in ALL_COMPONENTS}[c] for c in entity["components"]]
 95        assert "_type" in entity["data"], f"_type not in serialized data for {entity} (id: {entity['entity_id']}"
 96
 97        if entity_data["_type"]["type"] == "cube":
 98            texture_path = get_project_root() / tp if (tp := entity_data["_args"]["texture_path"]) is not None else None
 99            eid = make_cube(world, pose=np.float32(entity_data["pose"]), scale=entity_data["_args"]["scale"],
100                            size=np.float32(entity_data["_args"]["size"]), texture_path=texture_path)
101            if GroundFloor in entity_components:
102                world.update()
103                world.get_entity(eid).add_component(GroundFloor)
104
105        elif entity_data["_type"]["type"] == "mesh":
106            texture_path = get_project_root() / tp if (tp := entity_data["_args"]["texture_path"]) is not None else None
107            collidable = HasCollision in entity_components
108            make_mesh(world, pose=np.float32(entity_data["pose"]), scale=entity_data["_args"]["scale"],
109                      model_path=get_project_root() / entity_data["_args"]["model_path"], texture_path=texture_path,
110                      collidable=collidable)
111
112        elif entity_data["_type"]["type"] == "robot":
113            kwargs = {}
114            if HasPhysicsLevel1 in entity_components:
115                kwargs = {
116                    "physics_type": "level1",
117                    "max_velocities": np.float32(entity_data["max_velocities"]),
118                    "velocity": np.float32(entity_data["velocity"]),
119                    "acceleration": np.float32(entity_data["acceleration"]),
120                }
121            elif HasPhysicsLevel2 in entity_components:
122                kwargs = {
123                    "physics_type": "level2",
124                    "drag_coefficient": np.float32(entity_data["drag_coefficient"]),
125                    "max_accelerations": np.float32(entity_data["max_accelerations"]),
126                    "velocity": np.float32(entity_data["velocity"]),
127                    "acceleration": np.float32(entity_data["acceleration"]),
128                }
129
130            make_robot(world, pose=np.float32(entity_data["pose"]),
131                       scale=entity_data["_args"]["scale"],
132                       model_path=get_project_root() / entity_data["_args"]["model_path"],
133                       **kwargs)
134        else:
135            raise NotImplementedError(entity_data["_type"]["type"])
136    # TODO: we should assert here that all the newly added entities have the same components as these that were in the
137    # previous state. This means we need to do world.update() and validate each added component.
138    return world
139
140# map stuff
141
142def build_default_scene_objects() -> World:
143    """builds the default basic scene objects for the simulator: n robots + a few obstacles"""
144    # NOTE: if plugins add their own components, ideally they should be in a separate plugin-owned world! Not enforced..
145    world = World(components=ALL_COMPONENTS, extra_metadata=["serializable", "comment"])
146
147    crate_png = RESOURCES_PATH / "textures/crate.png"
148    grass_png = RESOURCES_PATH / "textures/grass.png"
149
150    # Ground plane
151    eid = make_cube(world, pte(make_arr(0, 0, 0, 0, 0, 0)), make_arr(20, 0.1, 20), scale=1, texture_path=grass_png)
152    world.update()
153    world.get_entity(eid).add_component(GroundFloor)
154    # Some scattered cubes at ground level
155    make_cube(world, pte(make_arr(5.0, 0.5, 5.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
156    make_cube(world, pte(make_arr(5.0, 0.5, -3.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
157    make_cube(world, pte(make_arr(-4.0, 0.5, 2.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
158    make_cube(world, pte(make_arr(-6.0, 0.5, -5.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
159    make_cube(world, pte(make_arr(8.0, 0.5, 0.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
160    # A small tower (stacked cubes)
161    make_cube(world, pte(make_arr(-2.0, 0.5, -2.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
162    make_cube(world, pte(make_arr(-2.0, 1.5, -2.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
163    make_cube(world, pte(make_arr(-2.0, 2.5, -2.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
164    # Floating obstacles (for UAV to navigate around)
165    make_cube(world, pte(make_arr(3.0, 3.0, 3.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
166    make_cube(world, pte(make_arr(-3.0, 4.0, 0.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
167    make_cube(world, pte(make_arr(0.0, 5.0, -4.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
168    # Directly in UAV's POV at start (UAV at y=6, looking +Z)
169    make_cube(world, pte(make_arr(0.0, 5.5, 8.0, 0, 0, 0)), make_arr(1.0, 1.0, 1.0), scale=1, texture_path=crate_png)
170    # House model
171    make_mesh(world, RESOURCES_PATH / "models/house/house.obj", pte(make_arr(5, 2, 0, 0, 0, 0)),
172               scale=0.1, collidable=False)
173    return world
174
175def build_robots(world: World, n_robots: int, robot_models: list[str]):
176    """build the robots in the scene"""
177    _, c = get_closest_square(n_robots)
178    model_paths = [RESOURCES_PATH / "models" / x / f"{x}.obj" for x in robot_models]
179    scales = [0.8] * len(model_paths)
180
181    # physics_type = "level1"
182    # kwargs = {
183    #     "max_velocities": np.float32(LVL1_MAX_VELOCITIES),
184    # }
185    physics_type = "level2"
186    kwargs = {
187        "max_accelerations": np.float32(LVL2_MAX_ACCELERATIONS),
188        "drag_coefficient": np.float32((LVL2_DRAG_COEFF, )),
189    }
190
191    for i in range(n_robots):
192        pose = pose_from_position_target_up(position=make_arr(i % c, 6, i // c),
193                                            target=make_arr(i % c, 6, i // c + 1), up=make_arr(0, 1, 0))
194        model_path = model_paths[i % len(model_paths)]
195        scale = scales[i % len(model_paths)]
196        make_robot(world, model_path, pose, scale, physics_type, **kwargs)
197
198# Simulator singleton and state object
199
200@dataclass(kw_only=True)
201class SimState(Serializable):
202    """The state of the scene including the cameras, robot and so on"""
203    world_camera: Camera | None = None
204    topdown_camera: Camera | None = None
205    active_camera_type: str = "world" # Which camera is active based on it's type
206    fpv_ix: int = 0
207    wireframe_mode: bool = False
208    collision_render_mode: bool = False
209    collision_cell_size: Point3D = field(default_factory=lambda: make_arr(30, 30, 30))
210    display_uav_trace: bool = False
211    uav_traces: dict[EntityId, FixedSizeDict[tuple[float, float, float], bool]] = field(default_factory=dict)
212
213    def __post_init__(self):
214        if self.world_camera is None:
215            self.world_camera = Camera(make_arr(15, 12, 15), make_arr(0, 2, 0), make_arr(0, 1, 0),
216                                       fovy=45, projection=rl.CAMERA_PERSPECTIVE)
217        if self.topdown_camera is None:
218            self.topdown_camera = Camera(make_arr(0, 30, 0), make_arr(0, 0, 0), make_arr(0, 0, -1),
219                                         fovy=20, projection=rl.CAMERA_ORTHOGRAPHIC)
220
221    @property
222    def active_camera_label(self) -> str:
223        """returns a string representation of the active camera for HUD"""
224        match self.active_camera_type:
225            case "world": return "WORLD CAMERA (F1)"
226            case "topdown": return "TOP-DOWN CAMERA (F2)"
227            case "fpv": return "FPV CAMERA (F3)"
228            case _: raise ValueError(self.active_camera_type)
229
230    @overrides
231    def to_dict(self) -> dict:
232        return {
233            "world_camera": self.world_camera.to_dict(),
234            "topdown_camera": self.topdown_camera.to_dict(),
235            "active_camera_type": self.active_camera_type,
236            "fpv_ix": self.fpv_ix,
237            "wireframe_mode": self.wireframe_mode,
238            "collision_render_mode": self.collision_render_mode,
239            "collision_cell_size": self.collision_cell_size.tolist(),
240            "display_uav_trace": self.display_uav_trace,
241            "uav_traces": FixedSizeDict({k: list(v.keys()) for k, v in self.uav_traces.items()},
242                                        maxlen=UAV_TRACES_MAX_LEN),
243        }
244
245    @staticmethod
246    @overrides
247    def from_dict(state: dict) -> SimState:
248        return SimState(
249            world_camera=Camera.from_dict(state["world_camera"]),
250            topdown_camera=Camera.from_dict(state["topdown_camera"]),
251            active_camera_type=state["active_camera_type"],
252            fpv_ix=state["fpv_ix"],
253            wireframe_mode=state["wireframe_mode"],
254            collision_render_mode=state["collision_render_mode"],
255            collision_cell_size=np.float32(state["collision_cell_size"]),
256            display_uav_trace=state["display_uav_trace"],
257            uav_traces={int(eid): FixedSizeDict({tuple(pt): True for pt in points}, maxlen=UAV_TRACES_MAX_LEN)
258                        for eid, points in state["uav_traces"].items()},
259        )
260
261class Simulator(Restorable):
262    """The simulator object. Contains the scene (objects/robots), state and events. Wires to/from_dict as well."""
263    def __init__(self, world: World, state: SimState, connection_manager: ConnectionManager,
264                 plugins: PluginsManager, protocol: Protocol):
265        self.world = world
266        self.state = state
267        self.connection_manager = connection_manager
268        self.plugins = plugins
269        self.protocol = protocol
270
271        self.physics_ticks_stats = deque(maxlen=TICK_STATS_LEN)
272        self.render_ticks_stats = deque(maxlen=TICK_STATS_LEN)
273
274        self.world.update()
275        self.robot_eids = self.world.query(HasFPV).entity_ids.copy()
276        self.channel_id_to_robot_eid: dict[int, int] = {}
277        for i, (channel, eid) in enumerate(zip(self.channels, self.robot_eids)): # network channels are not part of ECS
278            self.world.get_entity(eid).add_component(HasChannel, channel=np.array([channel], "object"))
279            self.channel_id_to_robot_eid[i] = eid
280
281        self.world.update()
282
283    @property
284    def channels(self) -> list[ObjToSimChannel]:
285        """The communication channels that this simulator has"""
286        return self.connection_manager.channels
287
288    def render_main_camera(self, extra_drawables: list[Drawable] | None):
289        """renders the main camera (the raylib UI) of the simulator"""
290        rl.BeginDrawing()
291        rl.ClearBackground(DEFAULT_BACKGROUND)
292
293        rl.BeginMode3D(self.active_camera.camera[0])
294
295        if self.state.wireframe_mode:
296            rl.rlEnableWireMode()
297
298        if self.state.collision_render_mode:
299            self._render_collision_mode(extra_drawables)
300        else:
301            entity_id = self.robot_eids[self.state.fpv_ix].item() if self.state.active_camera_type == "fpv" else None
302            self._render(entity_id=entity_id, extra_drawables=extra_drawables)
303
304        if self.state.display_uav_trace:
305            for traces in self.state.uav_traces.values():
306                for trace in traces.keys():
307                    rl.DrawSphere(trace, 0.05, rl.GREEN)
308
309        rl.EndMode3D()
310
311        if self.state.wireframe_mode:
312            rl.rlDisableWireMode()
313
314        n_robots = len(self.world.query(HasFPV))
315        msg = f"{self.state.active_camera_label}\nObjects: {len(self.world)} (robots: {n_robots})"
316        rl.DrawText(msg.encode(), 10, 10, 20, rl.DARKGRAY)
317        rl.DrawFPS(rl.GetScreenWidth() - 80, 10)
318        rl.EndDrawing()
319
320    def render_fpv_cameras(self, extra_drawables: list[Drawable] | None):
321        """Renders each robot's FPV camera but only if it's streaming. Calls self.draw() for scene content."""
322        qr = self.world.query(HasChannel, HasFPV)
323        for i, (channel, fpv_camera, fpv_texture, fpv_data) in enumerate(zip(qr.channel, qr.fpv_camera,
324                                                                             qr.fpv_texture, qr.fpv_data)):
325            if channel.item().client is None:
326                continue
327            fpv_texture: "rl.RenderTexture" = fpv_texture.item()
328            fpv_data: FPVData = fpv_data.item()
329            camera: Camera = fpv_camera.item()
330
331            rl.BeginTextureMode(fpv_texture)
332            rl.ClearBackground(DEFAULT_BACKGROUND)
333            rl.BeginMode3D(camera.camera[0])
334            self._render(entity_id=qr.entity_ids[i].item(), extra_drawables=extra_drawables)
335            rl.EndMode3D()
336            rl.EndTextureMode()
337
338            fpv_img = rl.LoadImageFromTexture(fpv_texture.texture)
339            fpv_img_ptr = rl.ffi.new("Image *", fpv_img)
340            rl.ImageFlipVertical(fpv_img_ptr)
341            with fpv_data.lock:
342                fpv_data.frame = bytes(rl.ffi.buffer(fpv_img_ptr.data, len(fpv_data.frame)))
343                fpv_data.frame_id += 1
344                fpv_data._frame_compressed = None # pylint: disable=protected-access
345            rl.UnloadImage(fpv_img_ptr[0])
346
347    @property
348    def active_camera(self) -> Camera:
349        """The current active camera out of the 3: world, uav or topdown"""
350        match self.state.active_camera_type:
351            case "world": return self.state.world_camera
352            case "topdown": return self.state.topdown_camera
353            case "fpv": return self.world.get_entity(
354                self.world.query(HasFPV).entity_ids[self.state.fpv_ix]).fpv_camera[0]
355            case _: raise ValueError(self.state.active_camera_type)
356
357    @overrides
358    def to_dict(self) -> dict:
359        return {
360            "world": world_to_dict(self.world),
361            "state": self.state.to_dict(),
362            "plugins": self.plugins.to_dict(),
363        }
364
365    @overrides
366    def load_state_dict(self, state: dict):
367        state = deepcopy(state) # this is mostly for INIT_STATE as the code below mutates it while loading
368
369        self._load_world_from_dict(state["world"])
370        self.state = SimState.from_dict(state["state"])
371
372        old_names = set(self.plugins.plugin_names)
373        for dropped in old_names - state["plugins"].keys():
374            logger.warning(f"Plugin '{dropped}' was live but absent from loaded state. Dropping.")
375        self.plugins = PluginsManager.from_dict(state["plugins"])
376
377    def _load_world_from_dict(self, world_state: dict):
378        world = world_from_dict(world_state)
379        world.update()
380
381        self.robot_eids = world.query(HasFPV).entity_ids.copy()
382        assert (A:=len(self.robot_eids)) == (B:=len(self.channels)), f"Loaded robots: {A} doesn't match existing: {B}"
383        for i, (channel, eid) in enumerate(zip(self.channels, self.robot_eids)): # network channels must be updated too
384            world.get_entity(eid).add_component(HasChannel, channel=np.array([channel], "object"))
385            self.channel_id_to_robot_eid[i] = eid
386        world.update()
387
388        logger.info(f"Loaded world from state: {world}")
389        self.world = world
390
391    def _render(self, entity_id: int | None, extra_drawables: list[Drawable] | None): #noqa
392        """Renders the scene and all its objects + robots. Called from within a rl.BeginMode3D() block."""
393        qr = self.world.query(HasModel, HasPose)
394        for i, (model, pose, scale) in enumerate(zip(qr.model, qr.pose, qr.scale)):
395            if entity_id is not None and entity_id == qr.entity_ids[i]: # if in FPV mode then don't draw yourself
396                continue
397            rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
398
399        for drawable in (extra_drawables or []): # e.g. plugins
400            drawable.draw(self.world, entity_id)
401
402    def _render_collision_mode(self, extra_drawables: list[Drawable] | None):
403        """render the scene in 'collision mode' which means we don't render the meshes of collidables, but the shapes"""
404        qr = self.world.query(HasModel, HasPose, exclude=[HasCollision])
405        for model, pose, scale in zip(qr.model, qr.pose, qr.scale):
406            rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
407
408        qr = self.world.query(HasModel, HasPose, HasCollision, exclude=[GroundFloor])
409        collider_size = qr.collider_bbox[:, 1] - qr.collider_bbox[:, 0]
410        for pose, kind, radius, size, colliding in zip(qr.pose, qr.collider_kind, qr.collider_radii,
411                                                       collider_size, qr.is_colliding):
412            color = rl.RED if colliding else rl.GREEN
413            if kind == ColliderKinds.SPHERE:
414                rl.DrawSphere(pose[0:3, 3].tolist(), radius.item(), color)
415            elif kind == ColliderKinds.AABB:
416                rl.DrawCube(pose[0:3, 3].tolist(), *size.tolist(), color)
417            else:
418                raise NotImplementedError(kind)
419
420        qr = self.world.query(HasModel, HasPose, GroundFloor)
421        for model, pose, scale in zip(qr.model, qr.pose, qr.scale):
422            rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
423
424        for drawable in (extra_drawables or []): # e.g. plugins
425            drawable.draw(self.world, None)
426
427    def __repr__(self):
428        return f"[Simulator]\n{self.world}\n{self.plugins}\n{self.connection_manager}\n{self.protocol}"
429
430# I/O stuff: network/keyboard handlers
431
432def keyboard_handler(sim: Simulator):
433    """handle keyboard input for the main UI. We create a message in the same protocol of robosim-ncat/global hanlder"""
434    msg = None
435
436    # camera management
437    if rl.IsKeyPressed(rl.KEY_F1):
438        msg = {"cmd": "sim_set_camera", "type": "world", "id": 0} # dummy ids, needed for protocol
439    if rl.IsKeyPressed(rl.KEY_F2):
440        msg = {"cmd": "sim_set_camera", "type": "topdown", "id": 0} # dummy ids, needed for protocol
441    if rl.IsKeyPressed(rl.KEY_F3): # we still keep the 'fpv' camera here but cannot control it.
442        if sim.state.active_camera_type == "fpv":
443            sim.state.fpv_ix = (sim.state.fpv_ix + 1) % len(sim.world.query(HasFPV))
444        msg = {"cmd": "sim_set_camera", "type": "fpv", "id": sim.state.fpv_ix}
445
446    # sim state management
447    if rl.IsKeyPressed(rl.KEY_F5): # save the state of the scene on the server side
448        msg = {"cmd": "sim_save_state"}
449    if rl.IsKeyPressed(rl.KEY_F6): # load the state of the scene on the server side
450        try:
451            with open(RESOURCES_PATH / "state.json", "r") as fp:
452                state = json.load(fp)
453            msg = {"cmd": "sim_load_state", "state": state}
454        except Exception as e:
455            logger.error(str(e))
456    if rl.IsKeyPressed(rl.KEY_R):
457        msg = {"cmd": "sim_load_state", "state": INIT_CONFIG}
458
459    # render management
460    if rl.IsKeyPressed(rl.KEY_I):
461        msg = {"cmd": "sim_set_wireframe_mode", "value": not sim.state.wireframe_mode}
462    if rl.IsKeyPressed(rl.KEY_O):
463        msg = {"cmd": "sim_set_collision_render_mode", "value": not sim.state.collision_render_mode}
464
465    # move robot via keyboard
466    if sim.state.active_camera_type == "fpv":
467        entity_id = sim.robot_eids[sim.state.fpv_ix]
468        control_input = np.float32([
469            rl.IsKeyDown(rl.KEY_A)       - rl.IsKeyDown(rl.KEY_D),           # left/right (x)
470            rl.IsKeyDown(rl.KEY_PAGE_UP) - rl.IsKeyDown(rl.KEY_PAGE_DOWN),   # up/down (y)
471            rl.IsKeyDown(rl.KEY_W)       - rl.IsKeyDown(rl.KEY_S),           # forward/backward (z)
472            rl.IsKeyDown(rl.KEY_DOWN)    - rl.IsKeyDown(rl.KEY_UP),          # pitch
473            rl.IsKeyDown(rl.KEY_Q)       - rl.IsKeyDown(rl.KEY_E),           # yaw
474            rl.IsKeyDown(rl.KEY_RIGHT)   - rl.IsKeyDown(rl.KEY_LEFT),        # roll
475        ])
476        if (control_input != 0).any():
477            msg = {"cmd": "entity_move", "entity_id": entity_id.item(), "control_input": control_input}
478
479    # misc / other commands
480    if rl.IsKeyPressed(rl.KEY_F9):
481        msg = {"cmd": "sim_set_collision_cell_size", "value": np.maximum(sim.state.collision_cell_size - 0.1, 0.1)}
482    if rl.IsKeyPressed(rl.KEY_F10):
483        msg = {"cmd": "sim_set_collision_cell_size", "value": np.minimum(sim.state.collision_cell_size + 0.1, 30)}
484    if rl.IsKeyPressed(rl.KEY_T):
485        msg = {"cmd": "sim_display_uav_trace", "value": not sim.state.display_uav_trace}
486
487    if msg is not None:
488        res = _handle_msg(sim, msg)
489        if "error" in res:
490            logger.error(res["error"])
491        else:
492            if msg["cmd"] != "entity_move":
493                logger.debug(res["status"])
494
495def network_handler(msg: dict, global_channel: ObjToSimChannel, connected_channel: ObjToSimChannel | None,
496                    sim: Simulator) -> dict | None:
497    """Ran in TCP thread. Fast path for one-off messages. If not fast -> sends to the right channel (global/conn).
498    Returns None if the message requires to be authenticated ('connect'), but the client isn't (channel is None)"""
499    assert "cmd" in msg, msg # should be handlded in connection manager already
500
501    all_cmds = [*list(sim.protocol.endpoints), *sim.plugins.all_endpoints]
502    if (cmd := msg["cmd"]) not in all_cmds:
503        return {"error": f"Unknown command: '{cmd}'", "supported_commands": all_cmds}
504
505    # these will go in the plugins manager and get handled by the right plugin
506    if cmd in sim.plugins.all_endpoints:
507        # If the client is not connected (e.g. controls a robot), then the handler returns and client gets disconnected
508        # TODO: it may be the case that some plugin messages should work even non connected (e.g. mission_get_state)
509        if connected_channel is None:
510            return None
511
512        with connected_channel.lock:
513            connected_channel.tcp2sim.put(msg)
514            return connected_channel.sim2tcp.get(timeout=SOCKET_TIMEOUT_S)
515
516    # guaranteed to be an admin/global command here
517
518    if (err := sim.protocol.validate_endpoint(endpoint=cmd, data=msg)) is not None:
519        logger.error(err)
520        return {"error": err.error}
521
522    if cmd == "help":
523        return {"supported_commands": all_cmds}
524    if cmd == "sim_get_info":
525        return {"control_loop_rate_hz": HZ,
526                "physics_tick_s": {"window": len(pts := sim.physics_ticks_stats), "p50": np.median(pts).item()},
527                "render_tick_s": {"window": len(rts := sim.render_ticks_stats), "p50": np.median(rts).item()},}
528    if cmd == "sim_get_state":
529        return sim.to_dict()
530
531    if cmd in ("robot_get_state", "robot_get_state_with_camera"):
532        channel = sim.channels[robot_ix := msg["robot_ix"]]
533        robot_eid = sim.channel_id_to_robot_eid[robot_ix]
534        robot_entity = sim.world.get_entity(robot_eid)
535        robot_state = robot_entity.to_dict(serialization_field="serializable")
536        robot_state["connected"] = channel.client is not None
537        fpv_data: FPVData = robot_entity.fpv_data.item()
538
539        if cmd == "robot_get_state":
540            return {"robot": robot_state}
541
542        if robot_state["connected"] is False:
543            return {"error": f"Robot '{robot_ix}' is not connected. Cannot get FPV data"}
544
545        with fpv_data.lock:
546            res = {"robot": robot_state, "fpv_compressed": (fc := fpv_data.frame_compressed),
547                   "fpv_shape": (fs := fpv_data.frame_shape), "fpv_frame_id": fpv_data.frame_id}
548            logger.log_every_s(f"FPV: {len(fc)} bytes (cr {len(fc)/(np.prod(fs)*3/4)*100:.2f}%)", "DEBUG", True) # RGB
549            return res
550
551    # These need to run on the main thread, so we use a special channel
552    with global_channel.lock:
553        global_channel.tcp2sim.put(msg)
554        return global_channel.sim2tcp.get(timeout=SOCKET_TIMEOUT_S)
555
556def _handle_load_state(sim: Simulator, new_state: dict) -> dict:
557    try:
558        for plugin in sim.plugins:
559            if isinstance(plugin, TrajectoryMissionPlugin) and plugin.state.phase == "running":
560                return {"error": "cannot reset while mission is running"}
561        sim.load_state_dict(new_state)
562        return {"status": "Loaded state"}
563    except Exception as e:
564        logger.error(f"Error: {e}\n{''.join(traceback.format_exception(e))}")
565        return {"error": str(e)}
566
567def _handle_msg(sim: Simulator, msg: dict) -> dict:
568    if not isinstance(msg, dict) or "cmd" not in msg:
569        return {"error": f"Wrong message format: {msg}"}
570    cmd = msg["cmd"]
571
572    if (err := sim.protocol.validate_endpoint(endpoint=cmd, data=msg)) is not None:
573        logger.error(err)
574        return {"error": err.error}
575
576    if cmd == "sim_set_camera":
577        if msg["type"] == "fpv":
578            sim.state.fpv_ix = msg["id"]
579        sim.state.active_camera_type = msg["type"]
580        return {"status": f"changed camera to '{sim.state.active_camera_type}'"}
581
582    elif cmd == "map_load":
583        raise NotImplementedError
584
585    # sim state management
586
587    elif cmd == "sim_reset":
588        # TODO: validate loaded state
589        return _handle_load_state(sim, new_state=INIT_CONFIG)
590
591    elif cmd == "sim_load_state":
592        # TODO: validate loaded state
593        return _handle_load_state(sim, new_state=msg["state"])
594
595    elif cmd == "sim_save_state":
596        try:
597            with open(RESOURCES_PATH / "state.json", "w") as fp:
598                json.dump(sim.to_dict(), fp, indent=4)
599        except Exception as e:
600            return {"error": str(e)}
601        return {"status": "Stored the state of the simulator to 'resources/state.json'"}
602
603    # render management
604
605    elif cmd == "sim_set_wireframe_mode":
606        sim.state.wireframe_mode = msg["value"]
607        return {"status": f"Set wireframe mode to: {sim.state.wireframe_mode}"}
608
609    elif cmd == "sim_set_collision_render_mode":
610        sim.state.collision_render_mode = msg["value"]
611        return {"status": f"Set collision render mode to: {sim.state.collision_render_mode}"}
612
613    # entity control management (e.g. move entities or in the future add/remove components or set values)
614
615    elif cmd == "entity_move":
616        if (entity_id := msg["entity_id"]) not in sim.world.live_entities:
617            return {"error": f"Entity with id: {entity_id} not in the world"}
618
619        entity = sim.world.get_entity(entity_id)
620        if not entity.has_component(HasMotionInput):
621            return {"error": f"Entity with id: {entity_id} does not have 'HasMotionInput' component"}
622        try:
623            entity.motion_input[:] = msg["control_input"]
624        except Exception as e:
625            return {"error": str(e)}
626        return {"status": f"Entity id: {entity_id}. Applied control_input: {msg['control_input']}"}
627
628    # misc / other commands
629
630    elif cmd == "sim_set_collision_cell_size":
631        sim.state.collision_cell_size = msg["value"]
632        return {"status": f"Set collision cell size to: {msg['value']}"}
633
634    elif cmd == "sim_display_uav_trace":
635        sim.state.uav_traces.clear()
636        sim.state.display_uav_trace = msg["value"]
637        return {"status": f"Display UAV trace set to to: {msg['value']}"}
638
639    raise ValueError(f"Shouldn't reach here: {msg}")
640
641def global_handler(sim: Simulator):
642    """Ran in main thread. Handles the admin commands that do not run in the 'fast path' & need main thread (opengl)"""
643    try:
644        msg = sim.connection_manager.global_channel.tcp2sim.get_nowait()
645    except Empty:
646        return
647    sim.connection_manager.global_channel.sim2tcp.put(_handle_msg(sim, msg))
648
649# main and cli args
650
651def main(args: Namespace):
652    """main fn"""
653    global INIT_CONFIG # pylint: disable=global-statement
654    rl.SetConfigFlags(rl.FLAG_WINDOW_HIDDEN | rl.FLAG_WINDOW_UNDECORATED if args.headless else rl.FLAG_WINDOW_RESIZABLE)
655    rl.InitWindow(SCREEN_WIDTH, SCREEN_HEIGHT, b"UAV Trajectory Simulation")
656    rl.SetWindowPosition(20, 200)
657    rl.SetTargetFPS(0)
658    rl.DisableEventWaiting()
659    rl.rlEnableBackfaceCulling()
660
661    if args.map_path is None:
662        world = build_default_scene_objects()
663    else:
664        with open(args.map_path, "r") as fp:
665            world = world_from_dict(json.load(fp)["world"])
666    build_robots(world, args.n_robots, args.robot_models)
667
668    state = SimState(collision_cell_size=make_collision_cell_size(world))
669    conn_manager = ConnectionManager("0.0.0.0", args.port, n_channels=args.n_robots, connect_command=CONNECT_CMD)
670
671    plugins = PluginsManager(plugin_names=args.plugins)
672    with open(args.protocol_path, "r") as fp:
673        protocol = Protocol.from_dict(json.load(fp), max_robot_ix=args.n_robots - 1)
674    sim = Simulator(world=world, state=state, connection_manager=conn_manager, plugins=plugins, protocol=protocol)
675    logger.info("Created default map" if args.map_path is None else f"Constructed map from '{args.map_path}'")
676    logger.info(f"Created {len(world.query(HasModel, exclude=[HasFPV]))} objects and {len(world.query(HasFPV))} robots")
677    logger.info(f"Created simulator object:\n{sim}")
678    logger.info(f"Raylib rendering on: {' / '.join(rl_get_device_and_renderer())}")
679    INIT_CONFIG = sim.to_dict()
680    conn_manager.start(partial(network_handler, sim=sim))
681
682    fpv_system = FPVCameraSystem()
683    physics_system = PhysicsSystem()
684
685    clock = Clock(dt=DT, max_ticks=MAX_SUBTICKS_PER_RENDER_TICK)
686    while True:
687        sim.world.update()
688        clock.wait_and_tick()
689
690        # I/O systems: read messages from the keyboard as well as all robots (via plugins/channels)
691        if rl.IsKeyPressed(rl.KEY_ESCAPE) or not conn_manager.tcp_thread.is_alive():
692            break
693        sim.world.query(HasMotionInput).motion_input[:] = 0 # reset all the inputs for the physics system
694        keyboard_handler(sim)
695        global_handler(sim)
696        sim.plugins.io_handler(world=sim.world)
697
698        # Update systems
699        if not args.headless:
700            sim.active_camera.manual_movement()
701        fpv_system(world=sim.world)
702
703        if sim.state.display_uav_trace:
704            robots = sim.world.query(HasFPV, HasPose)
705            for robot_id, pose in zip(robots.entity_ids, robots.pose):
706                robot_traces = sim.state.uav_traces.setdefault(int(robot_id), FixedSizeDict(maxlen=UAV_TRACES_MAX_LEN))
707                robot_traces[tuple(pose[0:3, 3].tolist())] = True
708
709        # Physics systems + the plugins callbacks before/after it
710        prev = time.perf_counter()
711        sim.plugins.on_before_physics(sim.world)
712        for _ in clock.drain():
713            physics_system(world=sim.world, dt=DT, cell_size=sim.state.collision_cell_size)
714        sim.plugins.on_after_physics(sim.world)
715        sim.physics_ticks_stats.append(time.perf_counter() - prev)
716
717        # Rendering: first fpv cameras (for each robot), then the main UI camera
718        prev = time.perf_counter()
719        sim.render_fpv_cameras(extra_drawables=sim.plugins)
720        if not args.headless:
721            sim.render_main_camera(extra_drawables=sim.plugins.plugins)
722        sim.render_ticks_stats.append(time.perf_counter() - prev)
723
724        sim.plugins.responses_handler(world=sim.world)
725
726    rl.CloseWindow()
727
728if __name__ == "__main__":
729    parser = ArgumentParser()
730    parser.add_argument("--map_path", type=Path, help="Path to the map config (scene objects, not robots). "
731                                                      "If not set, manual scene is built")
732    parser.add_argument("--headless", action="store_true", help="If set, raylib is started in headless mode")
733    parser.add_argument("--port", "-p", type=int, help="The port for the TCP listener", default=42069)
734    parser.add_argument("--n_robots", type=int, default=2, help="The maximum number of robots (connectable) allowed")
735    parser.add_argument("--robot_models", nargs="+", help="A list of models for robots. If not set, uses all available")
736    parser.add_argument("--plugins", nargs="*", default=["manual_move"],
737                        help=f"Which plugins to start, if any. Registered: {list(PluginsManager.REGISTERED_PLUGINS)}")
738    parser.add_argument("--protocol_path", type=Path, default=get_project_root() / "server/protocol.json")
739    arg = parser.parse_args()
740    assert len(set(arg.plugins)) == len(arg.plugins), f"duplicates in --plugins: {arg.plugins}"
741    if arg.robot_models is None:
742        arg.robot_models = [x.name for x in (RESOURCES_PATH / "models").iterdir()
743                             if x.name.startswith("drone") and x.name != "drone_cube_head"]
744        logger.info(f"--robot_models not provided, using all found: {arg.robot_models}")
745    assert arg.protocol_path.exists(), f"Protocol file (protocol.json) not found at: '{arg.protocol_path}'"
746    main(arg)