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
  9from micronetcode import Message, SourceType
 10
 11from robolib.constants import SOCKET_TIMEOUT_S, HZ
 12from robolib.utils import FPVData, logger, ecs_data_as_np, get_project_root
 13from robolib.entities import make_nonserializable_data, add_entity
 14from robolib.netcode import RobosimClientState, RobosimClient
 15from plugins.trajectory_mission_plugin import TrajectoryMissionPlugin
 16
 17sys.path.append(str(get_project_root() / "src/robosim/"))
 18from simulator_singleton import Simulator # pylint: disable=wrong-import-order, import-error, wrong-import-position
 19
 20RESOURCES_PATH = get_project_root() / "resources"
 21
 22def _fast_handler(msg: Message, sim: Simulator) -> dict:
 23    client: RobosimClient = msg.client
 24    if not isinstance(msg.data, dict) or (cmd := msg.data.get("cmd")) is None:
 25        return {"error": f"'cmd' not in message: {msg.data}"}
 26
 27    # Special command: streaming FPV camera is responded straight from the client thread asap
 28    if cmd == "robot_get_state_with_camera":
 29        if (err := sim.protocol.validate_endpoint(endpoint=cmd, data=msg.data)) is not None:
 30            return {"error": err.error}
 31
 32        robot_eid = sim.robot_eids[msg.data["robot_ix"]]
 33        if robot_eid not in sim.robot_eid_to_channel_ix:
 34            return {"error": f"Robot '{client.channel.idx}' is not connected. Cannot get FPV data"}
 35
 36        robot_entity = sim.world.get_entity(robot_eid)
 37        robot_state = robot_entity.to_dict(serialization_field="serializable")
 38        robot_state["connected"] = True
 39
 40        fpv_data: FPVData = robot_entity.fpv_data.item()
 41        with fpv_data.lock:
 42            res = {"robot": robot_state, "fpv_compressed": (fc := fpv_data.frame_compressed),
 43                   "fpv_shape": (fs := fpv_data.frame_shape), "fpv_frame_id": fpv_data.frame_id}
 44            logger.log_every_s(f"FPV: {len(fc)} bytes (cr {len(fc)/(np.prod(fs)*3/4)*100:.2f}%)", "DEBUG", True) # RGB
 45            return res
 46
 47    # If it's any other message, then the slow handler reads it. It can be a sim or a plugin message, both are handled.
 48    return msg.client.channel.put_then_get(msg, timeout=SOCKET_TIMEOUT_S)
 49
 50def network_handler(msg: Message, sim: Simulator) -> dict:
 51    """
 52    [TCP client thread]. All the client messages arrive here.
 53    The FPV command is 'fast handled' while all the others are pushed to the simulator (global/plugins).
 54    Note: stop making admin/global channels. We don't need it! Each robot has its own channel of communication.
 55    """
 56    timestamp_recv = datetime.now().isoformat()
 57    resp = _fast_handler(msg, sim)
 58    resp["timestamp_recv"] = timestamp_recv
 59    resp["timestamp_resp"] = datetime.now().isoformat()
 60    return resp
 61
 62def _handle_load_state(sim: Simulator, new_state: dict) -> dict:
 63    try:
 64        for plugin in sim.plugins_manager.plugins:
 65            # TODO(181): leaky abstraction, we need to make sure that some plugins are "unstoppable".
 66            if isinstance(plugin, TrajectoryMissionPlugin) and plugin.state.phase == "running":
 67                return {"error": "cannot reset while mission is running"}
 68        sim.load_state_dict(new_state)
 69        return {"status": "Loaded state"}
 70    except Exception as e:
 71        return {"error": f"Error: {e}\n{''.join(traceback.format_exception(e))}"}
 72
 73def _handle_message(sim: Simulator, msg: Message, init_config: dict) -> dict:
 74    """handles an admin message (either global or keyboard) that needs to run on the main thread"""
 75    client: RobosimClient = msg.client
 76    data = msg.data
 77    cmd = data["cmd"]
 78
 79    try:
 80        if (err := sim.protocol.validate_endpoint(endpoint=cmd, data=msg.data)) is not None:
 81            return {"error": err.error}
 82    except KeyError:
 83        return {"error": f"Unknown command: '{cmd}'. Supported commands: {list(sim.all_commands)}"}
 84
 85    if cmd == "help":
 86        return {"supported_commands": list(sim.all_commands)}
 87
 88    if cmd == "connect":
 89        if client.state == RobosimClientState.STAGED:
 90            robot_ix = sim.assign_channel_to_first_free_robot(client.channel.idx)
 91            if robot_ix is None:
 92                return {"error": "no robot is free, try again later"}
 93
 94            client.robot_eid = sim.robot_eids[robot_ix]
 95            client.state = RobosimClientState.CONNECTED
 96            return {"status": "connected", "robot_ix": robot_ix, "eid": client.robot_eid}
 97        else: # CONNECTED
 98            return {"error": "already connected"}
 99
100    if cmd == "robot_get_state":
101        robot_ix = data["robot_ix"]
102        robot_eid = sim.robot_eids[robot_ix]
103        robot_entity = sim.world.get_entity(robot_eid)
104        robot_state = robot_entity.to_dict(serialization_field="serializable")
105        robot_state["connected"] = robot_eid in sim.robot_eid_to_channel_ix # released at 'on_disconnect' time
106        return {"robot": robot_state}
107
108    if cmd == "sim_get_info":
109        return {"control_loop_rate_hz": HZ,
110                "physics_tick_s": {"window": len(pts := sim.physics_ticks_stats), "p50": np.median(pts).item()},
111                "render_tick_s": {"window": len(rts := sim.render_ticks_stats), "p50": np.median(rts).item()},}
112
113    if cmd == "sim_get_state":
114        return sim.to_dict()
115
116    if cmd == "sim_set_camera":
117        if data["type"] == "fpv":
118            sim.state.fpv_ix = data["id"]
119        sim.state.active_camera_type = data["type"]
120        return {"status": f"changed camera to '{sim.state.active_camera_type}'"}
121
122    elif cmd == "map_load":
123        raise NotImplementedError
124
125    # sim state management
126
127    elif cmd == "sim_reset":
128        # TODO: validate loaded state
129        return _handle_load_state(sim, new_state=init_config)
130
131    elif cmd == "sim_load_state":
132        # TODO: validate loaded state
133        return _handle_load_state(sim, new_state=data["state"])
134
135    elif cmd == "sim_save_state":
136        try:
137            with open(RESOURCES_PATH/"state.json", "w") as fp:
138                json.dump(sim.to_dict(), fp, indent=4)
139        except Exception as e:
140            return {"error": str(e)}
141        return {"status": "Stored the state of the simulator to 'resources/state.json'"}
142
143    # render management
144
145    elif cmd == "sim_set_wireframe_mode":
146        sim.state.wireframe_mode = data["value"]
147        return {"status": f"Set wireframe mode to: {sim.state.wireframe_mode}"}
148
149    elif cmd == "sim_set_collision_render_mode":
150        sim.state.collision_render_mode = data["value"]
151        return {"status": f"Set collision render mode to: {sim.state.collision_render_mode}"}
152
153    # entity control management (e.g. move entities or in the future add/remove components or set values)
154
155    elif cmd == "entity_spawn":
156        try:
157            components = [sim.world.component_name_to_type[c] for c in data["components"]]
158            eid = add_entity(sim.world, components, **data["data"])
159            sim.world.update()
160            return {"status": "Added a new entity to the world", "entity_id": eid}
161        except Exception as e:
162            return {"error": str(e)}
163
164    elif cmd == "entity_destroy":
165        try:
166            sim.world.remove_entity(entity_id := data["entity_id"])
167            sim.world.update()
168            return {"status": f"Entity {entity_id} destroyed"}
169        except Exception as e:
170            return {"error": str(e)}
171
172    elif cmd == "entity_add_component":
173        try:
174            entity = sim.world.get_entity(entity_id := data["entity_id"])
175            comp = sim.world.component_name_to_type[data["component"]]
176
177            field_dtype = dict(zip(sim.world.component_to_field_names[comp], sim.world.component_to_dtypes[comp]))
178            component_data = {k: ecs_data_as_np(v, dtype=field_dtype[k]) for k, v in data["data"].items()}
179
180            entity.add_component(comp, **component_data, **make_nonserializable_data([comp], component_data))
181            sim.world.update()
182            return {"status": f"Added component {data['component']} to entity {entity_id}"}
183        except Exception as e:
184            return {"error": str(e)}
185
186    elif cmd == "entity_remove_component":
187        try:
188            entity = sim.world.get_entity(entity_id := data["entity_id"])
189            component = sim.world.component_name_to_type[data["component"]]
190
191            entity.remove_component(component)
192            sim.world.update()
193            return {"status": f"Component {data['component']} removed from entity {entity_id}"}
194        except Exception as e:
195            return {"error": str(e)}
196
197    elif cmd == "entity_set_data":
198        try:
199            entity = sim.world.get_entity(entity_id := data["entity_id"])
200            comp = sim.world.component_name_to_type[data["component"]]
201            # TODO(microecs-30): can microecs autoconvert the data using numpy's rules?
202            field_dtype = dict(zip(sim.world.component_to_field_names[comp], sim.world.component_to_dtypes[comp]))
203            component_data = {k: ecs_data_as_np(v, dtype=field_dtype[k]) for k, v in data["data"].items()}
204
205            entity.set_data(**component_data)
206            sim.world.update()
207            return {"status": f"Entity id: {entity_id}. Updated data of following fields: {list(data['data'])}"}
208        except Exception as e:
209            return {"error": str(e)}
210
211    # misc / other commands
212
213    elif cmd == "sim_set_collision_cell_size":
214        sim.state.collision_cell_size = data["value"]
215        return {"status": f"Set collision cell size to: {data['value']}"}
216
217    elif cmd == "sim_display_uav_trace":
218        sim.state.uav_traces.clear()
219        sim.state.display_uav_trace = data["value"]
220        return {"status": f"Display UAV trace set to to: {data['value']}"}
221
222    raise ValueError(f"Shouldn't reach here: {msg}")
223
224def network_slow_handler(sim: Simulator, init_config: dict):
225    """[main thread] Polls the messages (global or plugin) of all the clients which were pushed from fast handler."""
226    messages_for_plugins: list[Message] = []
227
228    while True:
229        try:
230            msg = sim.connection_manager.get_one_message()
231        except Empty:
232            break
233
234        cmd = msg.data["cmd"] # Existence of 'cmd' was validated in _fast_handler
235
236        # 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
237        if cmd in sim.plugins_manager.all_endpoints:
238            messages_for_plugins.append(msg)
239            continue
240
241        # Below only messages for the core simulator (not plugins).
242
243        resp = _handle_message(sim, msg, init_config=init_config)
244        # Only respond if the messages are from a socket client (e.g. client is not None). Ignore if SCRIPT (keyboard).
245        if msg.source_type == SourceType.SOCKET:
246            msg.client.channel.main2tcp.put(resp)
247
248    # Call the plugins at the end. Note that this is needed even for zero messages as we call on_tick() over there.
249    # Maybe we need to split this? plugins.io_handler vs plugins.on_tick ?
250    sim.plugins_manager.io_handler(world=sim.world, messages_for_plugins=messages_for_plugins)