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 microspec import Protocol
 25from micronetcode import Message, ConnectionManager
 26
 27from robolib.utils import logger, get_project_root, FixedSizeDict, rl_get_device_and_renderer
 28from robolib.camera import Camera
 29from robolib.constants import SCREEN_HEIGHT, SCREEN_WIDTH, DT
 30from robolib.components import HasModel, HasPose, HasFPV, HasMotion
 31from robolib.systems import FPVRenderSystem, PhysicsSystem, RenderSystem
 32from robolib.netcode import MsgpackCodec, RobosimClient
 33from robolib.serialization import world_from_dict
 34
 35from plugins import REGISTERED_PLUGINS
 36
 37np.set_printoptions(precision=3, linewidth=120)
 38rl.SetTraceLogLevel(rl.LOG_WARNING)
 39
 40sys.path.append(str(get_project_root() / "src/robosim/"))
 41from simulator_singleton import Simulator, SimState # pylint: disable=wrong-import-order, import-error, wrong-import-position
 42from protocol import io_handler, network_handler, responses_handler, SCRIPT_CLIENT # pylint: disable=wrong-import-order, import-error, wrong-import-position
 43
 44# global settings and types
 45
 46PROTOCOL_PATH = get_project_root() / "src/robosim/protocol.json"
 47INIT_CONFIG: dict | None = None
 48
 49def _build_script_lines(script_path: Path) -> dict[int, list[Message]]:
 50    script_lines = {}
 51    with open(script_path, "r") as fp:
 52        for script_line in map(lambda x: x.strip(), fp.readlines()):
 53            if len(script_line) == 0 or script_line.startswith("#"): # comment or empty line
 54                continue
 55
 56            raw_timestamp, cmd, raw_args = script_line.split(" ", maxsplit=2)
 57            parsed_args = json.loads(raw_args)
 58            if "." in raw_timestamp or "-" in raw_timestamp:
 59                raise ValueError(f"Negative or floating point timestamp: {raw_timestamp}")
 60
 61            data = {"cmd": cmd, **parsed_args}
 62            script_lines.setdefault(int(raw_timestamp), []).append(Message("script", data, client=SCRIPT_CLIENT))
 63
 64    return script_lines
 65
 66def rl_init(headless: bool):
 67    """Initialize the window with proper flags, full screen (windowed) mode based on monitor width/height etc."""
 68    rl.SetConfigFlags(rl.FLAG_WINDOW_HIDDEN | rl.FLAG_WINDOW_UNDECORATED if headless else rl.FLAG_WINDOW_RESIZABLE)
 69    rl.InitWindow(SCREEN_WIDTH, SCREEN_HEIGHT, b"UAV Trajectory Simulation")
 70    rl.SetWindowPosition(20, 200)
 71    rl.SetTargetFPS(0)
 72    rl.DisableEventWaiting()
 73    rl.rlEnableBackfaceCulling()
 74    rl.SetExitKey(rl.KEY_NULL)
 75
 76def main(args: Namespace) -> list[dict]:
 77    """main fn"""
 78    global INIT_CONFIG # pylint: disable=global-statement
 79    rl_init(args.headless)
 80
 81    with open(args.map_path, "r") as fp:
 82        world = world_from_dict(json.load(fp)["world"])
 83    world.update()
 84
 85    state = SimState()
 86    conn_manager = ConnectionManager("0.0.0.0", args.port, codec=MsgpackCodec(), max_connections=args.max_connections,
 87                                     client_type=RobosimClient)
 88
 89    with open(PROTOCOL_PATH, "r") as fp:
 90        protocol = Protocol.from_dict(json.load(fp), max_connection_id=args.max_connections - 1,
 91                                      components=world.component_names)
 92
 93    script_lines, script_responses = {}, []
 94    if args.script is not None:
 95        script_lines = _build_script_lines(args.script)
 96        logger.info(f"Read {sum(len(x) for x in script_lines.values())} script lines from: '{args.script}'")
 97
 98    sim = Simulator(world=world, state=state, connection_manager=conn_manager, plugins=args.plugins, protocol=protocol)
 99    logger.info("Created default map" if args.map_path is None else f"Constructed map from '{args.map_path}'")
100    logger.info(f"Created {len(world.query(HasModel, exclude=[HasFPV]))} objects and {len(world.query(HasFPV))} robots")
101    logger.info(f"Created simulator object:\n{sim}")
102    logger.info(f"Raylib rendering on: {' / '.join(rl_get_device_and_renderer())}")
103    INIT_CONFIG = sim.to_dict()
104    if len(world.query(HasFPV)) == 0:
105        logger.warning("No robots found in the world. Add HasFPV component to existing entities")
106
107    conn_manager.on_disconnect = sim.client_on_disconnect
108    conn_manager.network_handler = partial(network_handler, sim=sim)
109    conn_manager.start()
110
111    fpv_render_system = FPVRenderSystem()
112    physics_system = PhysicsSystem()
113    render_system = RenderSystem()
114
115    cam: Camera
116    while True:
117        sim.world.update()
118        sim.clock.wait_and_tick()
119
120        # I/O systems: read messages from the keyboard as well as all robots (via plugins/channels)
121        if sim.state.must_exit or not conn_manager.is_alive():
122            break
123        sim.world.query(HasMotion).motion_input = 0 # reset all the inputs for the physics system
124        script_msgs_this_tick = script_lines.get(sim.clock.tick_count, [])
125        message_responses = io_handler(sim, extra_messages=script_msgs_this_tick, init_config=INIT_CONFIG)
126
127        # Update systems
128        if not args.headless:
129            sim.active_camera.manual_movement()
130        # Update FPV cameras based on current positions (TODO: should be somewhere else)
131        qr = sim.world.query(HasPose, HasFPV)
132        for cam, pose in zip(map(lambda x: x.item(), qr.fpv_camera), qr.pose):
133            cam.set_pose(position=pose[0:3, 3], up=pose[0:3, 1], target=(pose[0:3, 3] + pose[0:3, 2]))
134
135        if sim.state.display_uav_trace:
136            robots = sim.world.query(HasFPV, HasPose)
137            for robot_id, pose in zip(robots.entity_ids, robots.pose):
138                robot_traces = sim.state.uav_traces.setdefault(
139                    int(robot_id), FixedSizeDict(maxlen=sim.state.uav_traces_max_len))
140                robot_traces[tuple(pose[0:3, 3].tolist())] = True
141
142        # Physics systems + the plugins callbacks before/after it
143        prev = time.perf_counter()
144        sim.plugins_manager.on_before_physics(sim.world)
145        for _ in sim.clock.drain():
146            physics_system(world=sim.world, dt=DT, cell_size=sim.collision_cell_size)
147        sim.plugins_manager.on_after_physics(sim.world)
148        sim.physics_ticks_stats.append(time.perf_counter() - prev)
149
150        # Rendering: first fpv cameras (for each robot), then the main UI camera
151        rl.BeginDrawing()
152        prev = time.perf_counter()
153        camera_entity_id = sim.robot_eids[sim.state.fpv_ix] if sim.active_camera.label == "FPV" else None
154        if not args.headless:
155            render_system(sim.world, sim.active_camera, extra_drawables=sim.plugins_manager.plugins,
156                          camera_entity_id=camera_entity_id, collision_mode=sim.state.collision_render_mode,
157                          wireframe_mode=sim.state.wireframe_mode, display_uav_trace=sim.state.display_uav_trace,
158                          uav_traces=sim.state.uav_traces, display_hud=True)
159        fpv_render_system(world=sim.world, connected_eids=list(sim.robot_eid_to_channel_ix),
160                          extra_drawables=sim.plugins_manager.plugins)
161        sim.render_ticks_stats.append(time.perf_counter() - prev)
162        rl.EndDrawing()
163
164        responses_handler(sim, message_responses, script_responses)
165
166    rl.CloseWindow()
167    conn_manager.close()
168    return script_responses
169
170if __name__ == "__main__":
171    parser = ArgumentParser()
172    parser.add_argument("--map_path", type=Path, help="Path to the map config (scene objects, not robots). "
173                                                      "If not set, the default scene is loaded (maps/01-basic.json)")
174    parser.add_argument("--headless", action="store_true", help="If set, raylib is started in headless mode")
175    parser.add_argument("--port", "-p", type=int, help="The port for the TCP listener", default=42069)
176    parser.add_argument("--max_connections", type=int, default=2, help="The maximum number of robot connections")
177    parser.add_argument("--plugins", nargs="*", default=["manual_move"],
178                        help=f"Which plugins to start, if any. Registered: {list(REGISTERED_PLUGINS)}")
179    parser.add_argument("--script", "-i", type=Path, help="Optional path to a scrip to be ran at start time")
180    arg = parser.parse_args()
181    assert len(set(arg.plugins)) == len(arg.plugins), f"duplicates in --plugins: {arg.plugins}"
182    if arg.map_path is None:
183        arg.map_path = get_project_root() / "resources/maps/01-basic.json"
184        logger.info(f"No --map_path provided. Defaulting to: '{arg.map_path}'")
185    main(arg)