robosim.robosim

  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 pathlib import Path
 16from functools import partial
 17from argparse import ArgumentParser, Namespace
 18import json
 19import sys
 20import time
 21
 22import numpy as np
 23import raylib as rl
 24from microecs import World, EntityId
 25from microspec import Protocol
 26from micronetcode import Message, ConnectionManager
 27
 28from robolib.utils import (Clock, logger, get_project_root, get_closest_square,
 29                           pose_from_position_target_up, pose_from_trans_euler as pte, make_arr, FixedSizeDict,
 30                           rl_get_device_and_renderer)
 31from robolib.constants import (SCREEN_HEIGHT, SCREEN_WIDTH, DT, MAX_SUBTICKS_PER_RENDER_TICK,
 32                               LVL2_DRAG_COEFF, LVL2_MAX_ACCELERATIONS)
 33from robolib.physics import make_collision_cell_size
 34from robolib.entities import add_entity
 35from robolib.components import (ColliderKinds, HasModel, HasPose, HasFPV,
 36                                HasMotionInput, HasCollision, HasTag, HasPhysicsLevel1, HasPhysicsLevel2,
 37                                HasVelocity, HasAcceleration)
 38from robolib.systems import FPVCameraSystem, PhysicsSystem
 39from robolib.netcode import MsgpackCodec, RobosimClient
 40
 41from plugins import REGISTERED_PLUGINS
 42
 43np.set_printoptions(precision=3, linewidth=120)
 44rl.SetTraceLogLevel(rl.LOG_WARNING)
 45
 46sys.path.append(str(get_project_root() / "src/robosim/"))
 47from simulator_singleton import ALL_COMPONENTS, Simulator, SimState, world_from_dict, GroundFloor # pylint: disable=wrong-import-order, import-error, wrong-import-position
 48from protocol import network_slow_handler, network_handler, _handle_message # pylint: disable=wrong-import-order, import-error, wrong-import-position
 49
 50# global settings and types
 51
 52RESOURCES_PATH = get_project_root() / "resources"
 53INIT_CONFIG: dict | None = None
 54
 55# default map stuff
 56
 57def _make_cube(world: World, pose: np.ndarray, scale: float) -> EntityId:
 58    base_components = [HasTag, HasModel, HasPose]
 59    cube_obj = RESOURCES_PATH / "models/cube/cube.obj"
 60    crate_png = RESOURCES_PATH / "textures/crate.png"
 61    aabb = ColliderKinds.AABB
 62    cb_args = {"tag": "cube", "model_path": str(cube_obj), "texture_path": str(crate_png),
 63                "scale": [scale], "pose": pose, "collider_kind": [aabb]}
 64    return add_entity(world, base_components + [HasCollision], **cb_args)
 65
 66def build_default_scene_objects() -> World:
 67    """builds the default basic scene objects for the simulator: n robots + a few obstacles"""
 68    # NOTE: if plugins add their own components, ideally they should be in a separate plugin-owned world! Not enforced..
 69    world = World(components=ALL_COMPONENTS, extra_metadata=["serializable", "comment"])
 70
 71    grass_png = RESOURCES_PATH / "textures/grass.png"
 72    ground_obj = RESOURCES_PATH / "models/ground/ground.obj"
 73    house_obj = RESOURCES_PATH / "models/house/house.obj"
 74    aabb = ColliderKinds.AABB
 75    base_components = [HasTag, HasModel, HasPose]
 76    make_pte = lambda *data: pte(make_arr(*data))
 77
 78    # Ground plane
 79    gp_args = {"tag": "ground", "model_path": str(ground_obj), "texture_path": str(grass_png),
 80               "scale": [20.0], "pose": make_pte(0, 0, 0, 0, 0, 0), "collider_kind": [aabb]}
 81    eid = add_entity(world, components=base_components + [GroundFloor, HasCollision], **gp_args)
 82    world.update()
 83    ground_entity = world.get_entity(eid)
 84    height = (ground_entity.model[0].bbox[1][1] * ground_entity.scale).item()
 85    # Some scattered cubes at ground level (height (0.25) + 0.5 (box height//2) == ground level)
 86    _make_cube(world, make_pte(5.0,  height + 0.5,  5.0, 0, 0, 0), scale=1)
 87    _make_cube(world, make_pte(5.0,  height + 0.5, -3.0, 0, 0, 0), scale=1)
 88    _make_cube(world, make_pte(-4.0, height + 0.5,  2.0, 0, 0, 0), scale=1)
 89    _make_cube(world, make_pte(-6.0, height + 0.5, -5.0, 0, 0, 0), scale=1)
 90    _make_cube(world, make_pte(8.0,  height + 0.5,  0.0, 0, 0, 0), scale=1)
 91    # A small tower (stacked cubes)
 92    _make_cube(world, make_pte(-2.0,     height + 0.5, -2.0, 0, 0, 0), scale=1)
 93    _make_cube(world, make_pte(-2.0, 2 * height + 0.5, -2.0, 0, 0, 0), scale=1)
 94    _make_cube(world, make_pte(-2.0, 3 * height + 0.5, -2.0, 0, 0, 0), scale=1)
 95    # Floating obstacles (for UAV to navigate around)
 96    _make_cube(world, make_pte(3.0,  3.0,  3.0, 0, 0, 0), scale=1)
 97    _make_cube(world, make_pte(-3.0, 4.0,  0.0, 0, 0, 0), scale=1)
 98    _make_cube(world, make_pte(0.0,  5.0, -4.0, 0, 0, 0), scale=1)
 99    # Directly in UAV's POV at start (UAV at y=6, looking +Z)
100    _make_cube(world, make_pte(0.0, 5.5, 8.0, 0, 0, 0), scale=1)
101    # House model
102    house_args = {"tag": "house", "model_path": str(house_obj), "texture_path": None,
103                  "scale": [0.1], "pose": make_pte(5, 2, 0, 0, 0, 0)}
104    add_entity(world, components=base_components, **house_args)
105    return world
106
107def build_robots(world: World, n_robots: int, robot_models: list[str]):
108    """build the robots in the scene"""
109    _, c = get_closest_square(n_robots)
110    model_paths = [RESOURCES_PATH/ "models" / x / f"{x}.obj" for x in robot_models]
111    scales = [0.8] * len(model_paths)
112
113    # physics_type = "level1"
114    # kwargs = {
115    #     "max_velocities": np.float32(LVL1_MAX_VELOCITIES),
116    # }
117    physics_type = "level2"
118    physics_args = {
119        "max_accelerations": np.float32(LVL2_MAX_ACCELERATIONS),
120        "drag_coefficient": np.float32((LVL2_DRAG_COEFF, )),
121    }
122
123    for i in range(n_robots):
124        pose = pose_from_position_target_up(position=make_arr(i % c, 6, i // c),
125                                            target=make_arr(i % c, 6, i // c + 1), up=make_arr(0, 1, 0))
126        model_path = model_paths[i % len(model_paths)]
127        scale = scales[i % len(model_paths)]
128        components = [
129            HasTag, HasModel, HasPose, HasFPV,
130            HasPhysicsLevel1 if physics_type == "level1" else HasPhysicsLevel2,
131            HasVelocity, HasAcceleration, HasMotionInput, HasCollision,
132        ]
133
134        data = {
135            "tag": "robot",
136            "model_path": str(model_path),
137            "texture_path": None,
138            "scale": [scale],
139            "pose": pose,
140            "collider_kind": [ColliderKinds.SPHERE],
141            **physics_args,
142        }
143        add_entity(world, components, **data)
144
145def keyboard_handler(sim: Simulator):
146    """
147    Handle keyboard input for the main UI. We create a message in the same protocol of robosim-ncat/global hanlder
148    Note: network handler is in protocol.py definit the entire protocol handlers of protocol.json.
149    """
150    msgs: list[dict] = []
151
152    # camera management
153    if rl.IsKeyPressed(rl.KEY_F1):
154        msgs += [{"cmd": "sim_set_camera", "type": "world", "id": 0}] # dummy ids, needed for protocol
155    if rl.IsKeyPressed(rl.KEY_F2):
156        msgs += [{"cmd": "sim_set_camera", "type": "topdown", "id": 0}] # dummy ids, needed for protocol
157    if rl.IsKeyPressed(rl.KEY_F3): # we still keep the 'fpv' camera here but cannot control it.
158        if sim.state.active_camera_type == "fpv":
159            sim.state.fpv_ix = (sim.state.fpv_ix + 1) % len(sim.world.query(HasFPV))
160        msgs += [{"cmd": "sim_set_camera", "type": "fpv", "id": sim.state.fpv_ix}]
161
162    # sim state management
163    if rl.IsKeyPressed(rl.KEY_F5): # save the state of the scene on the server side
164        msgs += [{"cmd": "sim_save_state"}]
165    if rl.IsKeyPressed(rl.KEY_F6): # load the state of the scene on the server side
166        try:
167            with open(RESOURCES_PATH/"state.json", "r") as fp:
168                state = json.load(fp)
169            msgs += [{"cmd": "sim_load_state", "state": state}]
170        except Exception as e:
171            logger.error(str(e))
172    if rl.IsKeyPressed(rl.KEY_R):
173        msgs += [{"cmd": "sim_reset"}]
174
175    # render management
176    if rl.IsKeyPressed(rl.KEY_I):
177        msgs += [{"cmd": "sim_set_wireframe_mode", "value": not sim.state.wireframe_mode}]
178    if rl.IsKeyPressed(rl.KEY_O):
179        msgs += [{"cmd": "sim_set_collision_render_mode", "value": not sim.state.collision_render_mode}]
180
181    # move robot via keyboard
182    if sim.state.active_camera_type == "fpv":
183        entity_id = sim.robot_eids[sim.state.fpv_ix]
184        control_input = np.float32([
185            rl.IsKeyDown(rl.KEY_A)       - rl.IsKeyDown(rl.KEY_D),           # left/right (x)
186            rl.IsKeyDown(rl.KEY_PAGE_UP) - rl.IsKeyDown(rl.KEY_PAGE_DOWN),   # up/down (y)
187            rl.IsKeyDown(rl.KEY_W)       - rl.IsKeyDown(rl.KEY_S),           # forward/backward (z)
188            rl.IsKeyDown(rl.KEY_DOWN)    - rl.IsKeyDown(rl.KEY_UP),          # pitch
189            rl.IsKeyDown(rl.KEY_Q)       - rl.IsKeyDown(rl.KEY_E),           # yaw
190            rl.IsKeyDown(rl.KEY_RIGHT)   - rl.IsKeyDown(rl.KEY_LEFT),        # roll
191        ])
192        if (control_input != 0).any():
193            msgs += [{"cmd": "entity_set_data", "entity_id": entity_id, "component": "HasMotionInput",
194                     "data": {"motion_input": control_input}}]
195
196    # misc / other commands
197    if rl.IsKeyPressed(rl.KEY_F9):
198        msgs += [{"cmd": "sim_set_collision_cell_size", "value": np.maximum(sim.state.collision_cell_size - 0.1, 0.1)}]
199    if rl.IsKeyPressed(rl.KEY_F10):
200        msgs += [{"cmd": "sim_set_collision_cell_size", "value": np.minimum(sim.state.collision_cell_size + 0.1, 30)}]
201    if rl.IsKeyPressed(rl.KEY_T):
202        msgs += [{"cmd": "sim_display_uav_trace", "value": not sim.state.display_uav_trace}]
203
204    for msg in msgs:
205        res = _handle_message(sim, Message("script", data=msg), init_config=INIT_CONFIG)
206        if "error" in res:
207            logger.error(res["error"])
208        else:
209            if not (msg["cmd"] == "entity_set_data" and msg.get("component", "") == "HasMotionInput"):
210                logger.debug(res["status"])
211
212# main and cli args
213
214def main(args: Namespace):
215    """main fn"""
216    global INIT_CONFIG # pylint: disable=global-statement
217    rl.SetConfigFlags(rl.FLAG_WINDOW_HIDDEN | rl.FLAG_WINDOW_UNDECORATED if args.headless else rl.FLAG_WINDOW_RESIZABLE)
218    rl.InitWindow(SCREEN_WIDTH, SCREEN_HEIGHT, b"UAV Trajectory Simulation")
219    rl.SetWindowPosition(20, 200)
220    rl.SetTargetFPS(0)
221    rl.DisableEventWaiting()
222    rl.rlEnableBackfaceCulling()
223
224    if args.map_path is None:
225        world = build_default_scene_objects()
226    else:
227        with open(args.map_path, "r") as fp:
228            world = world_from_dict(json.load(fp)["world"])
229
230    world.update()
231    if len(world.query(HasFPV)) == 0:
232        build_robots(world, args.n_robots, args.robot_models)
233
234    state = SimState(collision_cell_size=make_collision_cell_size(world))
235    conn_manager = ConnectionManager("0.0.0.0", args.port, codec=MsgpackCodec(), max_connections=args.n_robots,
236                                     client_type=RobosimClient)
237
238    with open(args.protocol_path, "r") as fp:
239        protocol = Protocol.from_dict(json.load(fp), max_robot_ix=args.n_robots - 1, components=world.component_names)
240    sim = Simulator(world=world, state=state, connection_manager=conn_manager, plugins=args.plugins, protocol=protocol)
241    logger.info("Created default map" if args.map_path is None else f"Constructed map from '{args.map_path}'")
242    logger.info(f"Created {len(world.query(HasModel, exclude=[HasFPV]))} objects and {len(world.query(HasFPV))} robots")
243    logger.info(f"Created simulator object:\n{sim}")
244    logger.info(f"Raylib rendering on: {' / '.join(rl_get_device_and_renderer())}")
245    INIT_CONFIG = sim.to_dict()
246
247    conn_manager.on_disconnect = sim.client_on_disconnect
248    conn_manager.network_handler = partial(network_handler, sim=sim)
249    conn_manager.start()
250
251    fpv_system = FPVCameraSystem()
252    physics_system = PhysicsSystem()
253
254    clock = Clock(dt=DT, max_ticks=MAX_SUBTICKS_PER_RENDER_TICK)
255    while True:
256        sim.world.update()
257        clock.wait_and_tick()
258
259        # I/O systems: read messages from the keyboard as well as all robots (via plugins/channels)
260        if rl.IsKeyPressed(rl.KEY_ESCAPE) or not conn_manager.is_alive():
261            break
262        sim.world.query(HasMotionInput).motion_input = 0 # reset all the inputs for the physics system
263        keyboard_handler(sim)
264        network_slow_handler(sim, init_config=INIT_CONFIG)
265
266        # Update systems
267        if not args.headless:
268            sim.active_camera.manual_movement()
269        fpv_system(world=sim.world)
270
271        if sim.state.display_uav_trace:
272            robots = sim.world.query(HasFPV, HasPose)
273            for robot_id, pose in zip(robots.entity_ids, robots.pose):
274                robot_traces = sim.state.uav_traces.setdefault(
275                    int(robot_id), FixedSizeDict(maxlen=sim.state.uav_traces_max_len))
276                robot_traces[tuple(pose[0:3, 3].tolist())] = True
277
278        # Physics systems + the plugins callbacks before/after it
279        prev = time.perf_counter()
280        sim.plugins_manager.on_before_physics(sim.world)
281        for _ in clock.drain():
282            physics_system(world=sim.world, dt=DT, cell_size=sim.state.collision_cell_size)
283        sim.plugins_manager.on_after_physics(sim.world)
284        sim.physics_ticks_stats.append(time.perf_counter() - prev)
285
286        # Rendering: first fpv cameras (for each robot), then the main UI camera
287        prev = time.perf_counter()
288        sim.render_fpv_cameras(extra_drawables=sim.plugins_manager.plugins)
289        if not args.headless:
290            sim.render_main_camera(extra_drawables=sim.plugins_manager.plugins)
291        sim.render_ticks_stats.append(time.perf_counter() - prev)
292
293        sim.plugins_manager.responses_handler(world=sim.world)
294
295    rl.CloseWindow()
296
297if __name__ == "__main__":
298    parser = ArgumentParser()
299    parser.add_argument("--map_path", type=Path, help="Path to the map config (scene objects, not robots). "
300                                                      "If not set, manual scene is built")
301    parser.add_argument("--headless", action="store_true", help="If set, raylib is started in headless mode")
302    parser.add_argument("--port", "-p", type=int, help="The port for the TCP listener", default=42069)
303    parser.add_argument("--n_robots", type=int, default=2, help="The maximum number of robots (connectable) allowed")
304    parser.add_argument("--robot_models", nargs="+", help="A list of models for robots. If not set, uses all available")
305    parser.add_argument("--plugins", nargs="*", default=["manual_move"],
306                        help=f"Which plugins to start, if any. Registered: {list(REGISTERED_PLUGINS)}")
307    parser.add_argument("--protocol_path", type=Path, default=get_project_root() / "src/robosim/protocol.json")
308    arg = parser.parse_args()
309    assert len(set(arg.plugins)) == len(arg.plugins), f"duplicates in --plugins: {arg.plugins}"
310    if arg.robot_models is None:
311        arg.robot_models = [x.name for x in (RESOURCES_PATH/"models").iterdir()
312                             if x.name.startswith("drone") and x.name != "drone_cube_head"]
313        logger.info(f"--robot_models not provided, using all found: {arg.robot_models}")
314    assert arg.protocol_path.exists(), f"Protocol file (protocol.json) not found at: '{arg.protocol_path}'"
315    main(arg)