server.plugins.race_plugin

  1"""race_plugin.py - A plugin implementing a race between the UAV robots."""
  2from __future__ import annotations
  3import random
  4import os
  5from datetime import datetime
  6from dataclasses import dataclass
  7from enum import StrEnum
  8
  9from microecs import World
 10from microspec import Endpoint, Field, Dtype
 11
 12import numpy as np
 13import raylib as rl
 14from overrides import overrides
 15
 16from robosim.components import HasFPV, HasCollision, ColliderKinds
 17from robosim.utils import Point3D, logger, vec3_from, make_arr
 18from robosim.plugin import Plugin, Message, Response, Event
 19from robosim.traits import Serializable
 20from robosim.physics.collision import sphere_sphere_collision
 21
 22N_FLAGS = int(os.getenv("ROBOSIM_RACE_PLUGIN_N_FLAGS", "10"))
 23SEED = int(os.getenv("ROBOSIM_RACE_PLUGIN_SEED", "10"))
 24MAP_RANGE = [(-10, 10), (2, 20), (-10, 10)]
 25FLAG_RADIUS = 0.3
 26
 27def _rand(a: float, b: float) -> float:
 28    return random.random() * (b - a) + a
 29
 30class RacePhase(StrEnum):
 31    """The 'phase' of this race"""
 32    NOT_STARTED = "not_started"
 33    RUNNING     = "running"
 34    FINISHED    = "finished"
 35
 36@dataclass
 37class RaceState(Serializable):
 38    """the state of the race"""
 39    phase: RacePhase
 40    n_flags: int
 41    flags_positions: list[Point3D] | None = None
 42    n_robots: int | None = None
 43    robot_next_flag: np.ndarray | None = None
 44    start_ts: str | None = None
 45    end_ts: str | None = None
 46    winner: int | None = None
 47    _flags_collider: list[tuple[Point3D, float]] | None = None
 48
 49    def __post_init__(self):
 50        random.seed(SEED)
 51        if self.flags_positions is None:
 52            self.flags_positions = [make_arr(_rand(*MAP_RANGE[0]), _rand(*MAP_RANGE[1]), _rand(*MAP_RANGE[2]))
 53                                    for _ in range(self.n_flags)]
 54
 55    @property
 56    def flags_collider(self) -> list[tuple[Point3D, float]]:
 57        """returns the (position, radius) of each flag"""
 58        if self._flags_collider is None:
 59            self._flags_collider = [(position, FLAG_RADIUS) for position in self.flags_positions]
 60        return self._flags_collider
 61
 62    @overrides
 63    def to_dict(self) -> dict:
 64        return {
 65            "phase": self.phase,
 66            "n_flags": self.n_flags,
 67            "flags_positions": None if self.flags_positions is None else [x.tolist() for x in self.flags_positions],
 68            "n_robots": self.n_robots,
 69            "robot_next_flag": None if self.robot_next_flag is None else self.robot_next_flag.tolist(),
 70            "start_ts": self.start_ts,
 71            "end_ts": self.end_ts,
 72            "winner": self.winner,
 73        }
 74
 75    @staticmethod
 76    @overrides
 77    def from_dict(state: dict) -> RaceState:
 78        state["flags_positions"] = ([make_arr(*x) for x in state["flags_positions"]]
 79                                    if state["flags_positions"] is not None else None)
 80        state["robot_next_flag"] = (np.float32(state["robot_next_flag"])
 81                                    if state["robot_next_flag"] is not None else None)
 82        return RaceState(**state)
 83
 84class RacePlugin(Plugin):
 85    """RacePlugin implementation"""
 86    def __init__(self):
 87        self.state = RaceState(phase="not_started", n_flags=N_FLAGS)
 88
 89    @property
 90    @overrides
 91    def endpoints(self) -> list[Endpoint]:
 92        return [
 93            Endpoint("race_start", input={}, output={"status": Field(Dtype.STR)}),
 94            Endpoint("race_status", input={}, output={"status": Field(Dtype.DICT)}),
 95        ]
 96
 97    @overrides
 98    def on_message_receive(self, world: World, entity_id: int, message: Message) -> Response:
 99        if message.cmd == "race_start":
100            if self.state.phase == RacePhase.NOT_STARTED:
101                assert self.state.end_ts is None, self.state.to_dict()
102                self.state.n_robots = len(world.query(HasFPV))
103                self.state.robot_next_flag = np.zeros(self.state.n_robots, "uint32")
104                self.state.start_ts = datetime.now().isoformat()
105                self.state.phase = RacePhase.RUNNING
106                logger.info(f"Race started at {self.state.start_ts}. Message received from eid: {entity_id}")
107                return Response.ok("race starting")
108            elif self.state.phase == RacePhase.RUNNING:
109                return Response.err("race already started")
110            else: # race ended
111                return Response.err("race already ended")
112
113        if message.cmd == "race_status":
114            return Response.ok(self.state.to_dict())
115
116        raise ValueError("Shouldn't reach here")
117
118    @overrides
119    def on_message_response(self, world: World, event: Event) -> Response:
120        raise ValueError("Shouldn't reach here")
121
122    @overrides
123    def on_after_physics(self, world: World):
124        """callback called in the main loop after physics and before drawing (rendering)"""
125        if self.state.phase != RacePhase.RUNNING:
126            return
127
128        qr = world.query(HasFPV, HasCollision)
129        assert (qr.collider_kind.numpy() == ColliderKinds.SPHERE).all(), (qr.entity_ids, qr.collider_kind.numpy())
130
131        robots = [world.get_entity(id) for id in qr.entity_ids]
132        for i, robot in enumerate(robots):
133            assert robot.collider_kind.item() == ColliderKinds.SPHERE
134            # TODO: kinda hardcoded. We could make a world in this plugin and re-use check_collision on entities
135            flag_position, flag_radius = self.state.flags_collider[self.state.robot_next_flag[i]]
136            robot_position, robot_radius = robot.pose[0:3, 3], robot.collider_radii
137            if sphere_sphere_collision(flag_position, flag_radius, robot_position, robot_radius.item()):
138                self.state.robot_next_flag[i] += 1
139
140        if (self.state.robot_next_flag == self.state.n_flags).any(): # check if any robot reached the last flag
141            assert self.state.end_ts is None, self.state.to_dict()
142            self.state.end_ts = datetime.now().isoformat()
143            self.state.winner = np.where(self.state.robot_next_flag == self.state.n_flags)[0].item()
144            self.state.robot_next_flag = self.state.robot_next_flag * 0 + 1<<31
145            self.state.phase = RacePhase.FINISHED
146            logger.info(f"Robot {self.state.winner} (eid: {robot.entity_id}) won the race!")
147
148    @overrides
149    def draw(self, world: World, entity_id: int | None):
150        if entity_id is None or self.state.phase == RacePhase.NOT_STARTED: # not FPV-drawing (e.g. world camera)
151            self._draw_world_or_not_started(world)
152        else: # FPV camera and reace started
153            self._draw_fpv(world, entity_id)
154
155    @overrides
156    def to_dict(self) -> dict:
157        return self.state.to_dict()
158
159    @overrides
160    def load_state_dict(self, state: dict):
161        self.state = RaceState.from_dict(dict(state)) # dict() copy so it's not modified outside this call
162
163    def _draw_world_or_not_started(self, world: World):
164        for flag_ix, pos in enumerate(self.state.flags_positions):
165            rl.DrawSphere(vec3_from(pos), FLAG_RADIUS, rl.BLUE)
166            if flag_ix < len(self.state.flags_positions) - 1:
167                rl.DrawCylinderEx(vec3_from(pos), vec3_from(self.state.flags_positions[flag_ix + 1]),
168                                  0.01, 0.01, 5, rl.BLUE)
169
170    def _draw_fpv(self, world: World, entity_id: int):
171        robot_ix = list(world.query(HasFPV).entity_ids).index(entity_id)
172        for flag_ix, pos in enumerate(self.state.flags_positions):
173            # for FPV: draw green the next flag and red the others. Don't draw already captured flags.
174            is_next_flag = flag_ix == self.state.robot_next_flag[robot_ix]
175            color = rl.GREEN if is_next_flag else rl.RED
176            if flag_ix >= self.state.robot_next_flag[robot_ix]:
177                rl.DrawSphere(vec3_from(pos), FLAG_RADIUS, color)
178                # draw a line between current and next flag, but just not for the very next one.
179                if flag_ix < len(self.state.flags_positions) - 1:
180                    rl.DrawCylinderEx(vec3_from(pos), vec3_from(self.state.flags_positions[flag_ix + 1]),
181                                      0.01, 0.01, 5, rl.BLUE)