Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions config/level2.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@ check_race_track = true
check_drone_start_pos = true
# Lets you practice your controller without putting up gates & obstacles, assumes nominal positions given below.
real_track_objects = true
# Use the onboard Lighthouse positioning estimate instead of a motion capture system. Requires a
# Lighthouse deck + base stations, a single drone, and real_track_objects = false (no tracker can
# measure the gates/obstacles). Defaults to false (motion capture).
lighthouse = false

[[deploy.drones]]
id = 10
Expand Down
97 changes: 78 additions & 19 deletions lsy_drone_racing/envs/real_race_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ def __init__(
randomizations: ConfigDict,
sensor_range: float = 0.5,
control_mode: Literal["state", "attitude"] = "state",
lighthouse: bool = False,
):
"""Create a deployable version of the drone racing environment.

Expand All @@ -88,10 +89,19 @@ def __init__(
sensor_range: Sensor range. Determines at which distance the exact position of the
gates and obstacles is reveiled.
control_mode: Control mode of the drone.
lighthouse: If True, localize the drone with its onboard Lighthouse estimate read back
from the drone instead of a motion capture system (no ROS, single drone only).
"""
assert rclpy.ok(), "ROS2 is not running. Please start ROS2 before creating a deploy env."
self.lighthouse = lighthouse
if not lighthouse:
assert rclpy.ok(), "ROS2 is not running. Start it before creating a deploy env."
# Static env data
self.n_drones = len(drones)
if lighthouse and self.n_drones > 1:
raise NotImplementedError(
"Lighthouse mode is only supported for a single drone: each process can only read "
"its own drone's onboard estimate, with no shared tracker for the others."
)
self.gates, self.obstacles, self.drones = load_track(track)
self.n_gates = len(self.gates.pos)
self.n_obstacles = len(self.obstacles.pos)
Expand All @@ -113,8 +123,11 @@ def __init__(
radio_channel=drone_config["channel"],
drone_id=drone_config["id"],
drone_name=self.drone_name,
lighthouse=lighthouse,
)
self._ros_connector = (
None if lighthouse else ROSConnector(estimator_names=self.drone_names, timeout=10.0)
)
self._ros_connector = ROSConnector(estimator_names=self.drone_names, timeout=10.0)
# Dynamic data
self.data = EnvData.create(
n_drones=self.n_drones, n_gates=self.n_gates, n_obstacles=self.n_obstacles
Expand All @@ -127,6 +140,12 @@ def _reset(self, *, seed: int | None = None, options: dict | None = None) -> tup
# Update the position of gates and obstacles with the real positions measured from Mocap. If
# disabled, they are equal to the nominal positions defined in the track config.
if options.get("real_track_objects", True):
if self.lighthouse:
raise ValueError(
"real_track_objects=True is not supported in lighthouse mode (no tracker "
"to measure the gates/obstacles). Place the track at the configured "
"coordinates and set real_track_objects=False."
)
self._update_track_poses()
if options.get("check_race_track", True):
check_race_track(
Expand All @@ -138,17 +157,33 @@ def _reset(self, *, seed: int | None = None, options: dict | None = None) -> tup
nominal_obstacles_pos=self.obstacles.nominal_pos,
rng_config=self.randomizations,
)
if options.get("check_drone_start_pos", True):
check_drone_start_pos(
nominal_pos=self.drones.pos[self.rank],
real_pos=self._ros_connector.pos[self.drone_name],
rng_config=self.randomizations,
drone_name=self.drone_name,
)
self.data.reset(np.stack([self._ros_connector.pos[n] for n in self.drone_names]))

self.drone.connect(timeout=10.0)
self.drone.reset(arm=True)
if self.lighthouse:
# The onboard estimate is only available after the drone is connected and its estimator
# has converged, so bring it up first and arm only once the start position checks out.
self.drone.connect(timeout=10.0)
self.drone.reset(arm=False)
drone_pos, *_ = self._drone_states()
if options.get("check_drone_start_pos", True):
check_drone_start_pos(
nominal_pos=self.drones.pos[self.rank],
real_pos=drone_pos[self.rank],
rng_config=self.randomizations,
drone_name=self.drone_name,
)
self.data.reset(drone_pos)
self.drone.arm()
else:
if options.get("check_drone_start_pos", True):
check_drone_start_pos(
nominal_pos=self.drones.pos[self.rank],
real_pos=self._ros_connector.pos[self.drone_name],
rng_config=self.randomizations,
drone_name=self.drone_name,
)
self.data.reset(np.stack([self._ros_connector.pos[n] for n in self.drone_names]))
self.drone.connect(timeout=10.0)
self.drone.reset(arm=True)
self._last_drone_pos_update = 0 # Last time a position was sent to the drone estimator

return self.obs(), self.info()
Expand All @@ -164,9 +199,8 @@ def _step(self, action: NDArray) -> tuple[dict, float, bool, bool, dict]:
action[:3], action[3:6], action[6:9], action[9], action[10:]
)

drone_pos = np.stack([self._ros_connector.pos[drone] for drone in self.drone_names])
drone_pos, drone_quat, _, _ = self._drone_states()
assert drone_pos.dtype == np.float32, "Drone position must be of type float32"
drone_quat = np.stack([self._ros_connector.quat[drone] for drone in self.drone_names])
assert drone_quat.dtype == np.float32, "Drone quaternion must be of type float32"
# Check if the drone is in the sensor range of the gates and obstacles
dpos = drone_pos[:, None, :2] - self.gates.pos[None, :, :2]
Expand All @@ -192,6 +226,27 @@ def _step(self, action: NDArray) -> tuple[dict, float, bool, bool, dict]:
self._last_drone_pos_update = t
return self.obs(), self.reward(), self.terminated(), self.truncated(), self.info()

def _drone_states(self) -> tuple[NDArray, NDArray, NDArray, NDArray]:
"""Return stacked (pos, quat, vel, ang_vel) for all drones from the active source.

Uses the onboard Lighthouse estimate in lighthouse mode, otherwise the motion capture
system via ROS. All arrays are float32 to match the observation space.
"""
if self.lighthouse:
state = self.drone.get_obs()
return (
state["pos"][None, ...],
state["quat"][None, ...],
state["vel"][None, ...],
state["ang_vel"][None, ...],
)
rc = self._ros_connector
pos = np.stack([rc.pos[drone] for drone in self.drone_names])
quat = np.stack([rc.quat[drone] for drone in self.drone_names])
vel = np.stack([rc.vel[drone] for drone in self.drone_names])
ang_vel = np.stack([rc.ang_vel[drone] for drone in self.drone_names])
return pos, quat, vel, ang_vel

def obs(self) -> dict[str, NDArray]:
"""Return the observation of the environment."""
# If gates/obstacles are in sensor range use the actual pose, otherwise use the nominal pose
Expand All @@ -204,10 +259,7 @@ def obs(self) -> dict[str, NDArray]:
obstacles_pos = np.where(mask, self.obstacles.pos, self.obstacles.nominal_pos).astype(
np.float32
)
drone_pos = np.stack([self._ros_connector.pos[drone] for drone in self.drone_names])
drone_quat = np.stack([self._ros_connector.quat[drone] for drone in self.drone_names])
drone_vel = np.stack([self._ros_connector.vel[drone] for drone in self.drone_names])
drone_ang_vel = np.stack([self._ros_connector.ang_vel[drone] for drone in self.drone_names])
drone_pos, drone_quat, drone_vel, drone_ang_vel = self._drone_states()
obs = {
"pos": drone_pos,
"quat": drone_quat,
Expand Down Expand Up @@ -311,7 +363,8 @@ def close(self):
self.drone.close(emergency_stop=True)
finally:
# Close all ROS connections
self._ros_connector.close()
if self._ros_connector is not None:
self._ros_connector.close()


# region Single Drone Env
Expand Down Expand Up @@ -346,6 +399,7 @@ def __init__(
randomizations: ConfigDict,
sensor_range: float = 0.5,
control_mode: Literal["state", "attitude"] = "state",
lighthouse: bool = False,
):
"""Initialize the multi-drone environment.

Expand All @@ -371,6 +425,7 @@ def __init__(
sensor_range: Sensor range. Determines at which distance the exact position of the
gates and obstacles is reveiled.
control_mode: Control mode of the drone.
lighthouse: Use the onboard Lighthouse estimate instead of motion capture.
"""
super().__init__(
drones=drones,
Expand All @@ -380,6 +435,7 @@ def __init__(
randomizations=randomizations,
sensor_range=sensor_range,
control_mode=control_mode,
lighthouse=lighthouse,
)

def reset(self, *, seed: int | None = None, options: dict | None = None) -> tuple[dict, dict]:
Expand Down Expand Up @@ -450,6 +506,7 @@ def __init__(
randomizations: ConfigDict,
sensor_range: float = 0.5,
control_mode: Literal["state", "attitude"] = "state",
lighthouse: bool = False,
):
"""Initialize the multi-drone environment.

Expand All @@ -462,6 +519,7 @@ def __init__(
sensor_range: Sensor range. Determines at which distance the exact position of the
gates and obstacles is reveiled.
control_mode: Control mode of the drone.
lighthouse: Use the onboard Lighthouse estimate instead of motion capture.
"""
super().__init__(
drones=drones,
Expand All @@ -471,6 +529,7 @@ def __init__(
randomizations=randomizations,
sensor_range=sensor_range,
control_mode=control_mode,
lighthouse=lighthouse,
)

def reset(self, *, seed: int | None = None, options: dict | None = None) -> tuple[dict, dict]:
Expand Down
Loading