server.plugins.trajectory_mission_plugin

  1"""trajectory_mission_plugin.py code for trajectory-based mission (1st robot only)"""
  2from __future__ import annotations
  3from dataclasses import dataclass, field
  4from enum import StrEnum
  5
  6from overrides import overrides
  7from microecs import World, Entity
  8from microspec import Endpoint, Field, Dtype
  9import numpy as np
 10import raylib as rl
 11
 12from robosim.utils import Point3D, Pose4x4, logger, vec3_from
 13from robosim.components import HasFPV, HasCollision, HasPose
 14from robosim.constants import DT, N_DOF
 15from robosim.traits import Serializable
 16from robosim.plugin import Plugin, Message, Response, Event
 17
 18MISSION_WEIGHTS = (0.5, 0.4, 0.001)
 19
 20# pylint: disable=invalid-name
 21class MissionPhase(StrEnum):
 22    """The state of the mission"""
 23    NotStarted = "not-started"
 24    Running    = "running"
 25    Finished   = "finished"
 26
 27@dataclass
 28class MissionState(Serializable):
 29    """mission state fat struct"""
 30    phase: MissionPhase = MissionPhase.NotStarted
 31    via_points: list[Point3D] = field(default_factory=list) # The goal pose where the drone must get to (if any)
 32    # mode='running' parameters
 33    trajectory: np.ndarray | None = None # Must be set when state="running"
 34    trajectory_ix: int = 0 # Must be set when state="running"
 35    trajectory_poses: list[Pose4x4] = field(default_factory=list)
 36    # mode='done' parameters
 37    final_pose: Pose4x4 = None
 38    via_points_errors: np.ndarray = None
 39    via_closest_poses: list[Pose4x4] = field(default_factory=list)
 40    target_position_error: float = None
 41    total_steps: int = None
 42    final_error: float = None
 43    mission_errors: list[str] = field(default_factory=list)
 44
 45    @overrides
 46    def to_dict(self) -> dict:
 47        # Note: save only generic parameters (e.g. via_points) and finished (e.g. final_pose, error)
 48        return {
 49            "phase": self.phase,
 50            "via_points": [x.tolist() for x in self.via_points],
 51            "via_points_errors": self.via_points_errors.tolist() if self.via_points_errors is not None else None,
 52            "via_closest_poses": [x.tolist() for x in self.via_closest_poses],
 53            "final_pose": self.final_pose.tolist() if self.final_pose is not None else None,
 54            "target_position_error": self.target_position_error,
 55            "total_steps": self.total_steps,
 56            "final_error": self.final_error,
 57            "mission_errors": self.mission_errors,
 58        }
 59
 60    @staticmethod
 61    @overrides
 62    def from_dict(state: dict) -> Serializable:
 63        state["via_points"] = ([np.float32(x) for x in state["via_points"]]
 64                               if state["via_points"] is not None else None)
 65        state["via_points_errors"] = (np.float32(state["via_points_errors"])
 66                                      if state["via_points_errors"] is not None else None)
 67        state["via_closest_poses"] = ([np.float32(x) for x in state["via_closest_poses"]]
 68                                      if state["via_closest_poses"] is not None else None)
 69        state["final_pose"] = np.float32(state["final_pose"]) if state["final_pose"] is not None else None
 70        return MissionState(**state)
 71
 72    @property
 73    def target_pose(self) -> Pose4x4 | None:
 74        """returns the last via point, if any"""
 75        return self.via_points[-1] if len(self.via_points) > 0 else None
 76
 77    def start(self, trajectory: np.ndarray):
 78        """starts the mission given a trajectory"""
 79        assert self.phase in (MissionPhase.NotStarted, MissionPhase.Finished), self.phase
 80        # Important: reset all the mode='running' parameters (from old runs)
 81        self.trajectory = trajectory.astype("float32")
 82        self.trajectory_poses.clear()
 83        self.trajectory_ix = 0
 84        self.mission_errors = []
 85        self.phase = MissionPhase.Running
 86
 87    @property
 88    def motion_input(self) -> np.ndarray | None:
 89        """return the control input from the mission if it's running"""
 90        assert self.phase == MissionPhase.Running, self.to_dict()
 91        if self.trajectory_ix < len(self.trajectory):
 92            return self.trajectory[self.trajectory_ix]
 93        return None
 94
 95    def observe(self, world: World, event: Event):
 96        """The mission run method. Called on each frame once the mission is running."""
 97        if self.phase != MissionPhase.Running:
 98            return
 99
100        robot: Entity = world.get_entity(event.target)
101        if robot.is_colliding.item():
102            robot_state = robot.to_dict(serialization_field="serializable")
103            event_state = {"source": type(event.source).__name__,
104                           "payload": {k: v.tolist() for k, v in event.payload.items()}}
105            self.mission_errors.append(
106                f"Trajectory ix={self.trajectory_ix}. Event state: {event_state}. Robot state: {robot_state}")
107
108        self.trajectory_poses.append(robot.pose.copy())
109        self.trajectory_ix += 1
110
111        if len(self.trajectory) < 2:
112            logger.error(f"At least two trajectory poses needed: {self.trajectory_poses}. Cancelling mission.")
113            self.phase = MissionPhase.NotStarted
114            return
115
116        curr = self.trajectory_poses[-1]
117        prev = self.trajectory_poses[-2] if len(self.trajectory_poses) > 2 else curr
118        if self.trajectory_ix >= len(self.trajectory) and np.allclose(prev, curr, atol=1e-3, rtol=0):
119            # note: don't get rid of via points here as we may start another one
120            self.final_pose = robot.pose.copy()
121            err = np.linalg.norm(self.target_pose - robot.pose[0:3, 3])
122            self.target_position_error = err.item()
123
124            # compute the closest pose towards each via point
125            traj_positions = np.float32(self.trajectory_poses)[:, 0:3, 3] # (N, 3)
126            via_points = np.float32(self.via_points) # (M, 3)
127            distances = np.linalg.norm(traj_positions[:, None] - via_points[None], axis=2) # (M, N)
128            min_distances_ix = np.argmin(distances, axis=1) # (M, )
129            self.via_points_errors = distances[list(range(len(min_distances_ix))), min_distances_ix] # (M, )
130            self.via_closest_poses = [self.trajectory_poses[ix] for ix in min_distances_ix] # (M, 6)
131
132            self.total_steps = len(self.trajectory_poses)
133            self.final_error = ((self.target_position_error * MISSION_WEIGHTS[0] +
134                                 self.via_points_errors.mean() * MISSION_WEIGHTS[1] +
135                                 self.total_steps * DT * MISSION_WEIGHTS[2]) / sum(MISSION_WEIGHTS)).item()
136            logger.info(f"Mission ended. Final error: {self.final_error}")
137            self.phase = MissionPhase.Finished
138
139class TrajectoryMissionPlugin(Plugin):
140    """TrajectoryMissionPlugin implementation"""
141    def __init__(self):
142        self.state = MissionState()
143        self.events_this_tick: list[Event] = []
144        self.zeros = np.float32([0] * N_DOF)
145
146    @property
147    @overrides
148    def endpoints(self) -> list[Endpoint]:
149        return [
150            # TODO(microspec): add shape-range
151            Endpoint("mission_add_via_point",
152                     input={"position": Field(Dtype.FLOAT32, shape=(3, ), range=[-1000, 1000])},
153                     output={"status": Field(Dtype.STR)}),
154            Endpoint("mission_get_state", input={}, output={"status": Field(Dtype.DICT)}),
155            Endpoint("mission_start",
156                     input={"trajectory": Field(Dtype.FLOAT32, shape=(None, N_DOF), min_len=2)},
157                     output={"status": Field(Dtype.STR)}),
158            Endpoint("mission_clear_via_points", input={}, output={"status": Field(Dtype.STR)}),
159        ]
160
161    @overrides
162    def on_message_receive(self, world: World, entity_id: int, message: Message) -> Response:
163        robot = world.get_entity(world.query(HasFPV, HasCollision, HasPose).entity_ids[0])
164
165        if entity_id != robot.entity_id:
166            return Response.err(f"Only robot id 0 can do mission commands. You are robot id {entity_id}")
167
168        if message.cmd == "mission_get_state":
169            return Response.ok(self.state.to_dict())
170
171        if message.cmd == "mission_add_via_point": # add a new via point to the mission
172            self.state.via_points.append(message.payload["position"])
173            self.state.phase = MissionPhase.NotStarted # in case it's in finished mode
174            return Response.ok((f"Added a new via point at {message.payload["position"]} "
175                                f"(current: {len(self.state.via_points)})"))
176
177        if message.cmd == "mission_clear_via_points":
178            self.state.via_points.clear()
179            self.state.phase = MissionPhase.NotStarted # in case it's in finished mode
180            return Response.ok("Removed all via points")
181
182        if message.cmd == "mission_start":
183            if len(self.state.via_points) < 1: # Note: mission-specific when we have more missions
184                return Response.err(f"Too few points ({len(self.state.via_points)}) for mission. Need at least 1")
185
186            self.state.start(message.payload["trajectory"])
187            return Response.ok((f"Starting mission through {len(self.state.via_points)} via points and "
188                                f"{len(self.state.trajectory)} trajectory points"))
189
190        raise ValueError("shouldn't get here")
191
192    @overrides
193    def on_tick(self, world: World) -> list[Event]:
194        self.events_this_tick.clear()
195        if self.state.phase == MissionPhase.Running:
196            robot_id = world.query(HasFPV, HasCollision, HasPose).entity_ids[0]
197            payload = {"motion_input": self.zeros if self.state.motion_input is None else self.state.motion_input}
198            self.events_this_tick = [Event(source=self, target=robot_id, payload=payload)]
199        return self.events_this_tick
200
201    @overrides
202    def on_after_physics(self, world: World):
203        for event in self.events_this_tick:
204            if event.applied:
205                self.state.observe(world, event)
206
207    @overrides
208    def draw(self, world: World, entity_id: int | None = None):
209        for position in self.state.via_points:
210            rl.DrawSphere(vec3_from(position), 0.2, rl.GREEN)
211
212        # Draw trajectory path during/after mission
213        for pose in self.state.trajectory_poses:
214            pos = (pose[0, 3], pose[1, 3], pose[2, 3])
215            rl.DrawCube(pos, 0.05, 0.05, 0.05, rl.BLUE)
216
217    @overrides
218    def on_message_response(self, world: World, event: Event) -> Response:
219        raise ValueError(f"Should not reach here: ({event})")
220
221    @overrides
222    def to_dict(self) -> dict:
223        return self.state.to_dict()
224
225    @overrides
226    def load_state_dict(self, state: dict):
227        self.state = MissionState.from_dict(dict(state)) # dict() copy so it's not modified outside this call