robosim.protocol

  1"""protocol.py - implements all the robosim-related protocol handlers. New commands of protocol.json go here"""
  2import traceback
  3import json
  4import sys
  5from queue import Empty
  6from datetime import datetime
  7
  8import numpy as np
  9import raylib as rl
 10from micronetcode import Message, Channel, SourceType
 11
 12from robolib.components import HasFPV
 13from robolib.constants import SOCKET_TIMEOUT_S, HZ, RESOURCES_PATH
 14from robolib.utils import FPVData, logger, get_project_root
 15from robolib.serialization import add_entity, add_component
 16from robolib.netcode import RobosimClientState, RobosimClient
 17from robolib.plugin import Response, ResponseKind
 18from plugins.trajectory_mission_plugin import TrajectoryMissionPlugin
 19
 20sys.path.append(str(get_project_root() / "src/robosim/"))
 21from simulator_singleton import Simulator # pylint: disable=wrong-import-order, import-error, wrong-import-position
 22
 23SCRIPT_CLIENT = RobosimClient(socket=None, channel=Channel(idx=-69, channel_queue_maxsize=1))
 24KEYBOARD_CLIENT = RobosimClient(socket=None, channel=Channel(idx=-67, channel_queue_maxsize=1))
 25
 26# Public API
 27
 28def network_handler(msg: Message, sim: Simulator) -> dict:
 29    """
 30    [TCP client thread]. All the client messages arrive here.
 31    The FPV command is 'fast handled' while all the others are pushed to the simulator (global/plugins).
 32    Note: stop making admin/global channels. We don't need it! Each connection has its own channel of communication.
 33    """
 34    timestamp_recv = datetime.now().isoformat()
 35    if not isinstance(msg.data, dict) or (cmd := msg.data.get("cmd")) is None:
 36        resp = {"error": f"'cmd' not in message: {msg.data}"}
 37    # Special command: streaming FPV camera is responded straight from the client thread asap
 38    elif cmd == "robot_get_state_with_camera":
 39        resp = _get_state_with_camera(msg, sim)
 40    # If it's any other message, then the io_handler reads it below via get_one_message().
 41    # It can be a sim or a plugin message, both are handled the same.
 42    else:
 43        raw_resp: Response = msg.client.channel.put_then_get(msg, timeout=SOCKET_TIMEOUT_S)
 44        resp = raw_resp.payload
 45    resp["timestamp_recv"] = timestamp_recv
 46    resp["timestamp_resp"] = datetime.now().isoformat()
 47    return resp
 48
 49def io_handler(sim: Simulator, extra_messages: list[Message], init_config: dict) -> list[tuple[Message, Response]]:
 50    """
 51    [main thread] Polls the messages (global or plugin) of all the clients which were pushed from fast handler.
 52    Additional messages (extra_messages) can be passed from the caller. This is used by (tick-aware) scripting messages
 53    """
 54    if len(extra_messages) > 0:
 55        logger.debug(f"Handling {len(extra_messages)} extra script messages (t={sim.clock.tick_count})")
 56    keyboard_messages = _keyboard_handler(sim)
 57    messages_to_process: list[Message] = [*extra_messages, *keyboard_messages]
 58    messages_for_plugins: list[Message] = []
 59    responses: list[tuple[Message, Response]] = []
 60
 61    while True:
 62        try:
 63            messages_to_process.append(sim.connection_manager.get_one_message())
 64        except Empty:
 65            break
 66
 67    for msg in messages_to_process:
 68        cmd = msg.data["cmd"] # Existence of 'cmd' was validated in network_handler
 69
 70        # If it is a message for a plugin, then put it in this list and which is used by the plugins manager at the end
 71        if cmd in sim.plugins_manager.all_endpoints:
 72            messages_for_plugins.append(msg)
 73        else: # core messages are responded asap
 74            resp = _handle_message(sim, msg, init_config=init_config)
 75            kind = ResponseKind.Ok if "error" not in resp else ResponseKind.Error
 76            responses.append((msg, Response(kind, payload=resp)))
 77
 78    # Call the plugins at the end. Note that this is needed even for zero messages as we call on_tick() over there.
 79    plugin_responses = sim.plugins_manager.io_handler(world=sim.world, messages_for_plugins=messages_for_plugins)
 80    return responses + plugin_responses # these are passed by the main loop in responses_handler
 81
 82def responses_handler(sim: Simulator, message_responses: list[tuple[Message, Response]], script_responses: list):
 83    """[main thread] Handler that ensures that all the messages of this tick are responded to."""
 84    plugin_responses = sim.plugins_manager.responses_handler(world=sim.world)
 85    for msg, response in message_responses + plugin_responses:
 86        assert isinstance(response, Response), f"\n{msg}\n{response}"
 87        if msg.client is SCRIPT_CLIENT:
 88            script_responses.append(response.payload)
 89        elif msg.client is KEYBOARD_CLIENT:
 90            if "error" in response.payload:
 91                logger.error(response.payload["error"])
 92            else:
 93                if not (msg.data["cmd"] == "entity_set_data" and msg.data.get("component", "") == "HasMotion"):
 94                    logger.debug(response.payload["status"])
 95        else:
 96            msg.client.channel.main2tcp.put(response)
 97
 98# Private functions below
 99
100def _get_state_with_camera(msg: Message, sim: Simulator) -> dict:
101    if (err := sim.protocol.validate_endpoint(endpoint=msg.data["cmd"], data=msg.data)) is not None:
102        return {"error": err.error}
103
104    robot_eid = sim.robot_eids[msg.data["robot_ix"]]
105    if robot_eid not in sim.robot_eid_to_channel_ix:
106        return {"error": f"Robot '{msg.client.channel.idx}' is not connected. Cannot get FPV data"}
107
108    robot_entity = sim.world.get_entity(robot_eid)
109    robot_state = robot_entity.to_dict(serialization_field="serializable")
110    robot_state["connected"] = True
111
112    fpv_data: FPVData = robot_entity.fpv_data.item()
113    with fpv_data.lock:
114        res = {"robot": robot_state, "fpv_compressed": (fc := fpv_data.frame_compressed),
115                "fpv_shape": (fs := fpv_data.frame_shape), "fpv_frame_id": fpv_data.frame_id}
116        logger.log_every_s(f"FPV: {len(fc)} bytes (cr {len(fc)/(np.prod(fs)*3/4)*100:.2f}%)", "DEBUG", True) # RGB
117        return res
118
119def _handle_load_state(sim: Simulator, new_state: dict) -> dict:
120    try:
121        for plugin in sim.plugins_manager.plugins:
122            # TODO(181): leaky abstraction, we need to make sure that some plugins are "unstoppable".
123            if isinstance(plugin, TrajectoryMissionPlugin) and plugin.state.phase == "running":
124                return {"error": "cannot reset while mission is running"}
125        sim.load_state_dict(new_state)
126        return {"status": "Loaded state"}
127    except Exception as e:
128        return {"error": f"Error: {e}\n{''.join(traceback.format_exception(e))}"}
129
130def _keyboard_handler(sim: Simulator) -> list[Message]:
131    """
132    Handle keyboard input for the main UI. We create a message in the same protocol of robosim-ncat/global hanlder
133    Note: network handler is in protocol.py definit the entire protocol handlers of protocol.json.
134    """
135    res: list[dict] = []
136
137    if rl.IsKeyPressed(rl.KEY_ESCAPE):
138        res += [{"cmd": "sim_exit"}]
139
140    # camera management
141    if rl.IsKeyPressed(rl.KEY_F1):
142        res += [{"cmd": "sim_set_camera", "type": "world", "id": 0}] # dummy ids, needed for protocol
143    if rl.IsKeyPressed(rl.KEY_F2):
144        res += [{"cmd": "sim_set_camera", "type": "topdown", "id": 0}] # dummy ids, needed for protocol
145    if rl.IsKeyPressed(rl.KEY_F3): # we still keep the 'fpv' camera here but cannot control it.
146        if sim.state.active_camera_type == "fpv":
147            sim.state.fpv_ix = (sim.state.fpv_ix + 1) % len(sim.world.query(HasFPV))
148        res += [{"cmd": "sim_set_camera", "type": "fpv", "id": sim.state.fpv_ix}]
149
150    # sim state management
151    if rl.IsKeyPressed(rl.KEY_F5): # save the state of the scene on the server side
152        res += [{"cmd": "sim_save_state"}]
153    if rl.IsKeyPressed(rl.KEY_F6): # load the state of the scene on the server side
154        try:
155            with open(RESOURCES_PATH/"state.json", "r") as fp:
156                state = json.load(fp)
157            res += [{"cmd": "sim_load_state", "state": state}]
158        except Exception as e:
159            logger.error(str(e))
160    if rl.IsKeyPressed(rl.KEY_R):
161        res += [{"cmd": "sim_reset"}]
162
163    # render management
164    if rl.IsKeyPressed(rl.KEY_I):
165        res += [{"cmd": "sim_set_wireframe_mode", "value": not sim.state.wireframe_mode}]
166    if rl.IsKeyPressed(rl.KEY_O):
167        res += [{"cmd": "sim_set_collision_render_mode", "value": not sim.state.collision_render_mode}]
168
169    # move robot via keyboard
170    if sim.state.active_camera_type == "fpv":
171        entity_id = sim.robot_eids[sim.state.fpv_ix]
172        control_input = np.float32([
173            rl.IsKeyDown(rl.KEY_A)       - rl.IsKeyDown(rl.KEY_D),           # left/right (x)
174            rl.IsKeyDown(rl.KEY_PAGE_UP) - rl.IsKeyDown(rl.KEY_PAGE_DOWN),   # up/down (y)
175            rl.IsKeyDown(rl.KEY_W)       - rl.IsKeyDown(rl.KEY_S),           # forward/backward (z)
176            rl.IsKeyDown(rl.KEY_DOWN)    - rl.IsKeyDown(rl.KEY_UP),          # pitch
177            rl.IsKeyDown(rl.KEY_Q)       - rl.IsKeyDown(rl.KEY_E),           # yaw
178            rl.IsKeyDown(rl.KEY_RIGHT)   - rl.IsKeyDown(rl.KEY_LEFT),        # roll
179        ])
180        if (control_input != 0).any():
181            res += [{"cmd": "entity_set_data", "entity_id": entity_id, "component": "HasMotion",
182                     "data": {"motion_input": control_input}}]
183
184    # misc / other commands
185    if rl.IsKeyPressed(rl.KEY_F9):
186        res += [{"cmd": "sim_set_collision_cell_size", "value": np.maximum(sim.collision_cell_size - 0.1, 0.1)}]
187    if rl.IsKeyPressed(rl.KEY_F10):
188        res += [{"cmd": "sim_set_collision_cell_size", "value": np.minimum(sim.collision_cell_size + 0.1, 30)}]
189    if rl.IsKeyPressed(rl.KEY_T):
190        res += [{"cmd": "sim_display_uav_trace", "value": not sim.state.display_uav_trace}]
191
192    return [Message(source_type=SourceType.SCRIPT, data=item, client=KEYBOARD_CLIENT) for item in res]
193
194def _handle_message(sim: Simulator, msg: Message, init_config: dict) -> dict:
195    """handles an admin message (either global or keyboard) that needs to run on the main thread"""
196    client: RobosimClient = msg.client
197    data = msg.data
198    cmd = data["cmd"]
199
200    try:
201        if (err := sim.protocol.validate_endpoint(endpoint=cmd, data=msg.data)) is not None:
202            return {"error": err.error}
203    except KeyError:
204        return {"error": f"Unknown command: '{cmd}'. Supported commands: {sorted(sim.all_commands)}"}
205
206    if cmd == "help":
207        return {"supported_commands": sorted(sim.all_commands)}
208
209    if cmd == "connect":
210        if client.state == RobosimClientState.STAGED:
211            robot_ix = sim.assign_channel_to_first_free_robot(client.channel.idx)
212            if robot_ix is None:
213                return {"error": "no robot is free, try again later"}
214
215            client.robot_eid = sim.robot_eids[robot_ix]
216            client.state = RobosimClientState.CONNECTED
217            return {"status": "connected", "robot_ix": robot_ix, "eid": client.robot_eid}
218        else: # CONNECTED
219            return {"error": "already connected"}
220
221    if cmd == "disconnect":
222        if client.state == RobosimClientState.CONNECTED:
223            sim.connection_manager.on_disconnect(client)
224            return {"status": "disconnected"}
225        else:
226            return {"error": "client not connected"}
227
228    if cmd == "robot_get_state":
229        robot_ix = data["robot_ix"]
230        robot_eid = sim.robot_eids[robot_ix]
231        robot_entity = sim.world.get_entity(robot_eid)
232        robot_state = robot_entity.to_dict(serialization_field="serializable")
233        robot_state["connected"] = robot_eid in sim.robot_eid_to_channel_ix # released at 'on_disconnect' time
234        return {"robot": robot_state}
235
236    if cmd == "robot_get_state_with_camera": # this is normally responded via (fast) network_handler (except in tests)
237        return _get_state_with_camera(msg, sim)
238
239    if cmd == "sim_get_info":
240        return {"control_loop_rate_hz": HZ,
241                "physics_tick_s": {"window": len(pts := sim.physics_ticks_stats), "p50": np.median(pts).item()},
242                "render_tick_s": {"window": len(rts := sim.render_ticks_stats), "p50": np.median(rts).item()},}
243
244    if cmd == "sim_get_state":
245        return sim.to_dict()
246
247    if cmd == "sim_set_camera":
248        if data["type"] == "fpv":
249            sim.state.fpv_ix = data["id"]
250        sim.state.active_camera_type = data["type"]
251        return {"status": f"changed camera to '{sim.state.active_camera_type}'"}
252
253    elif cmd == "map_load":
254        raise NotImplementedError
255
256    elif cmd == "sim_exit":
257        sim.state.must_exit = True
258        return {"status": "Exitting"}
259
260    # sim state management
261
262    elif cmd == "sim_reset":
263        # TODO: validate loaded state
264        return _handle_load_state(sim, new_state=init_config)
265
266    elif cmd == "sim_load_state":
267        # TODO: validate loaded state
268        return _handle_load_state(sim, new_state=data["state"])
269
270    elif cmd == "sim_save_state":
271        try:
272            with open(RESOURCES_PATH/"state.json", "w") as fp:
273                json.dump(sim.to_dict(), fp, indent=4)
274        except Exception as e:
275            return {"error": str(e)}
276        return {"status": "Stored the state of the simulator to 'resources/state.json'"}
277
278    # render management
279
280    elif cmd == "sim_set_wireframe_mode":
281        sim.state.wireframe_mode = data["value"]
282        return {"status": f"Set wireframe mode to: {sim.state.wireframe_mode}"}
283
284    elif cmd == "sim_set_collision_render_mode":
285        sim.state.collision_render_mode = data["value"]
286        return {"status": f"Set collision render mode to: {sim.state.collision_render_mode}"}
287
288    # entity control management (e.g. move entities or in the future add/remove components or set values)
289
290    elif cmd == "entity_spawn":
291        try:
292            components = [sim.world.component_name_to_type[c] for c in data["components"]]
293            eid = add_entity(sim.world, components, **data["data"])
294            sim.world.update()
295            return {"status": "Added a new entity to the world", "entity_id": eid}
296        except Exception as e:
297            return {"error": str(e)}
298
299    elif cmd == "entity_destroy":
300        try:
301            sim.world.remove_entity(entity_id := data["entity_id"])
302            sim.world.update()
303            return {"status": f"Entity {entity_id} destroyed"}
304        except Exception as e:
305            return {"error": str(e)}
306
307    elif cmd == "entity_add_component":
308        try:
309            entity = sim.world.get_entity(entity_id := data["entity_id"])
310            comp = sim.world.component_name_to_type[data["component"]]
311            add_component(entity, component=comp, **data["data"])
312            sim.world.update()
313            return {"status": f"Added component {data['component']} to entity {entity_id}"}
314        except Exception as e:
315            return {"error": str(e)}
316
317    elif cmd == "entity_remove_component":
318        try:
319            entity = sim.world.get_entity(entity_id := data["entity_id"])
320            component = sim.world.component_name_to_type[data["component"]]
321            entity.remove_component(component)
322            sim.world.update()
323            return {"status": f"Component {data['component']} removed from entity {entity_id}"}
324        except Exception as e:
325            return {"error": str(e)}
326
327    elif cmd == "entity_set_data":
328        try:
329            entity = sim.world.get_entity(entity_id := data["entity_id"])
330            entity.set_data(**data["data"])
331            sim.world.update()
332            return {"status": f"Entity id: {entity_id}. Updated data of following fields: {list(data['data'])}"}
333        except Exception as e:
334            return {"error": str(e)}
335
336    # misc / other commands
337
338    elif cmd == "sim_set_collision_cell_size":
339        sim.collision_cell_size = data["value"]
340        return {"status": f"Set collision cell size to: {data['value']}"}
341
342    elif cmd == "sim_display_uav_trace":
343        sim.state.uav_traces.clear()
344        sim.state.display_uav_trace = data["value"]
345        return {"status": f"Display UAV trace set to to: {data['value']}"}
346
347    raise ValueError(f"Shouldn't reach here: {msg}")