Implement vectorized, multi-drone core race env.

Add specialized environments based off of RaceCoreEnv. Switch from rpy_rates to ang_vel

Author: amacati

benchmarks/sim.py

@@ -76,9 +76,9 @@
 import jax
 
 import lsy_drone_racing
+from lsy_drone_racing.envs.multi_drone_race import MultiDroneRaceEnv
 
 env = gymnasium.make('MultiDroneRacing-v0',
-    n_envs=1,  # TODO: Remove this for single-world envs
     n_drones=config.env.n_drones,
     freq=config.env.freq,
     sim_config=config.sim,
@@ -98,7 +98,7 @@
 # JIT masked reset (used in autoreset)
 mask = env.unwrapped.data.marked_for_reset
 mask = mask.at[0].set(True)
-env.unwrapped.reset(mask=mask)
+super(MultiDroneRaceEnv, env.unwrapped).reset(mask=mask)  # enforce masked reset compile
 jax.block_until_ready(env.unwrapped.data)
 env.action_space.seed(2)
 """

config/level0.toml

@@ -74,4 +74,4 @@ pos = [0.0, 1.0, 1.4]
 pos = [1.0, 1.0, 0.05]
 rpy = [0, 0, 0]
 vel = [0, 0, 0]
-rpy_rates = [0, 0, 0]
+ang_vel = [0, 0, 0]

config/level1.toml

@@ -61,7 +61,7 @@ pos = [0.0, 1.0, 1.4]
 pos = [1.0, 1.0, 0.05]
 rpy = [0, 0, 0]
 vel = [0, 0, 0]
-rpy_rates = [0, 0, 0]
+ang_vel = [0, 0, 0]
 
 [env.disturbances.action]
 fn = "normal"

config/level2.toml

@@ -61,7 +61,7 @@ pos = [0.0, 1.0, 1.4]
 pos = [1.0, 1.0, 0.05]
 rpy = [0, 0, 0]
 vel = [0, 0, 0]
-rpy_rates = [0, 0, 0]
+ang_vel = [0, 0, 0]
 
 [env.disturbances.action]
 fn = "normal"

config/level3.toml

@@ -61,7 +61,7 @@ pos = [-0.5, 0.0, 1.4]
 pos = [1.0, 1.0, 0.07]
 rpy = [0, 0, 0]
 vel = [0, 0, 0]
-rpy_rates = [0, 0, 0]
+ang_vel = [0, 0, 0]
 
 [env.disturbances.action]
 fn = "normal"

config/multi_level0.toml

@@ -66,4 +66,4 @@ pos = [0.0, 1.0, 1.4]
 pos = [[1.0, 1.0, 0.05], [1.2, 1.0, 0.05]]
 rpy = [[0, 0, 0], [0, 0, 0]]
 vel = [[0, 0, 0], [0, 0, 0]]
-rpy_rates = [[0, 0, 0], [0, 0, 0]]
+ang_vel = [[0, 0, 0], [0, 0, 0]]

config/multi_level3.toml

@@ -66,7 +66,7 @@ pos = [0.0, 1.0, 1.4]
 pos = [[1.0, 1.0, 0.05], [1.2, 1.0, 0.05]]
 rpy = [[0, 0, 0], [0, 0, 0]]
 vel = [[0, 0, 0], [0, 0, 0]]
-rpy_rates = [[0, 0, 0], [0, 0, 0]]
+ang_vel = [[0, 0, 0], [0, 0, 0]]
 
 [env.disturbances.action]
 fn = "normal"

lsy_drone_racing/envs/__init__.py

@@ -15,44 +15,46 @@
 
 from gymnasium import register
 
-from lsy_drone_racing.envs.drone_racing_env import DroneRacingEnv
-
-__all__ = ["DroneRacingEnv"]
+# region SingleDroneEnvs
 
 register(
     id="DroneRacing-v0",
-    entry_point="lsy_drone_racing.envs.drone_racing_env:DroneRacingEnv",
+    entry_point="lsy_drone_racing.envs.drone_race:DroneRaceEnv",
+    vector_entry_point="lsy_drone_racing.envs.drone_race:VecDroneRaceEnv",
     max_episode_steps=1800,  # 30 seconds * 60 Hz,
     disable_env_checker=True,  # Remove warnings about 2D observations
 )
 
 register(
     id="DroneRacingAttitude-v0",
-    entry_point="lsy_drone_racing.envs.drone_racing_env:DroneRacingAttitudeEnv",
+    entry_point="lsy_drone_racing.envs.drone_race:DroneRaceAttitudeEnv",
+    vector_entry_point="lsy_drone_racing.envs.drone_race:VecDroneRaceAttitudeEnv",
     max_episode_steps=1800,
     disable_env_checker=True,
 )
 
+# region MultiDroneEnvs
+
 register(
-    id="DroneRacingDeploy-v0",
-    entry_point="lsy_drone_racing.envs.drone_racing_deploy_env:DroneRacingDeployEnv",
+    id="MultiDroneRacing-v0",
+    entry_point="lsy_drone_racing.envs.multi_drone_race:MultiDroneRaceEnv",
+    vector_entry_point="lsy_drone_racing.envs.multi_drone_race:VecMultiDroneRaceEnv",
     max_episode_steps=1800,
     disable_env_checker=True,
 )
 
+# region DeployEnvs
+
 register(
-    id="DroneRacingAttitudeDeploy-v0",
-    entry_point="lsy_drone_racing.envs.drone_racing_deploy_env:DroneRacingAttitudeDeployEnv",
+    id="DroneRacingDeploy-v0",
+    entry_point="lsy_drone_racing.envs.drone_racing_deploy_env:DroneRacingDeployEnv",
     max_episode_steps=1800,
     disable_env_checker=True,
 )
 
-# region MultiEnvs
-
-# TODO: Register specialized, non-vectorized envs for single worlds
 register(
-    id="MultiDroneRacing-v0",
-    entry_point="lsy_drone_racing.envs.vec_drone_race:VectorMultiDroneRaceEnv",
+    id="DroneRacingAttitudeDeploy-v0",
+    entry_point="lsy_drone_racing.envs.drone_racing_deploy_env:DroneRacingAttitudeDeployEnv",
     max_episode_steps=1800,
     disable_env_checker=True,
 )

lsy_drone_racing/envs/drone_race.py

@@ -0,0 +1,110 @@
+"""Single drone racing environments."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from gymnasium import Env
+from gymnasium.vector import VectorEnv
+from gymnasium.vector.utils import batch_space
+
+from lsy_drone_racing.envs.race_core import RaceCoreEnv, action_space, observation_space
+
+if TYPE_CHECKING:
+    import numpy as np
+    from ml_collections import ConfigDict
+    from numpy.typing import NDArray
+
+
+class DroneRaceEnv(RaceCoreEnv, Env):
+    def __init__(
+        self,
+        freq: int,
+        sim_config: ConfigDict,
+        sensor_range: float,
+        track: ConfigDict | None = None,
+        disturbances: ConfigDict | None = None,
+        randomizations: ConfigDict | None = None,
+        random_resets: bool = False,
+        seed: int = 1337,
+        max_episode_steps: int = 1500,
+        device: Literal["cpu", "gpu"] = "cpu",
+    ):
+        super().__init__(
+            n_envs=1,
+            n_drones=1,
+            freq=freq,
+            sim_config=sim_config,
+            sensor_range=sensor_range,
+            track=track,
+            disturbances=disturbances,
+            randomizations=randomizations,
+            random_resets=random_resets,
+            seed=seed,
+            max_episode_steps=max_episode_steps,
+            device=device,
+        )
+        self.action_space = action_space("state")
+        n_gates, n_obstacles = len(track.gates), len(track.obstacles)
+        self.observation_space = observation_space(n_gates, n_obstacles)
+        self.autoreset = False
+
+    def reset(self, seed: int | None = None, options: dict | None = None) -> tuple[dict, dict]:
+        obs, info = super().reset(seed=seed, options=options)
+        obs = {k: v[0, 0] for k, v in obs.items()}
+        info = {k: v[0, 0] for k, v in info.items()}
+        return obs, info
+
+    def step(self, action: NDArray[np.floating]) -> tuple[dict, float, bool, bool, dict]:
+        obs, reward, terminated, truncated, info = super().step(action)
+        obs = {k: v[0, 0] for k, v in obs.items()}
+        info = {k: v[0, 0] for k, v in info.items()}
+        return obs, reward[0, 0], terminated[0, 0], truncated[0, 0], info
+
+
+class VecDroneRaceEnv(RaceCoreEnv, VectorEnv):
+    def __init__(
+        self,
+        num_envs: int,
+        freq: int,
+        sim_config: ConfigDict,
+        sensor_range: float,
+        track: ConfigDict | None = None,
+        disturbances: ConfigDict | None = None,
+        randomizations: ConfigDict | None = None,
+        random_resets: bool = False,
+        seed: int = 1337,
+        max_episode_steps: int = 1500,
+        device: Literal["cpu", "gpu"] = "cpu",
+    ):
+        super().__init__(
+            n_envs=num_envs,
+            n_drones=1,
+            freq=freq,
+            sim_config=sim_config,
+            sensor_range=sensor_range,
+            track=track,
+            disturbances=disturbances,
+            randomizations=randomizations,
+            random_resets=random_resets,
+            seed=seed,
+            max_episode_steps=max_episode_steps,
+            device=device,
+        )
+        self.single_action_space = action_space("state")
+        self.action_space = batch_space(self.single_action_space, num_envs)
+        n_gates, n_obstacles = len(track.gates), len(track.obstacles)
+        self.single_observation_space = observation_space(n_gates, n_obstacles)
+        self.observation_space = batch_space(self.single_observation_space, num_envs)
+
+    def reset(self, seed: int | None = None, options: dict | None = None) -> tuple[dict, dict]:
+        obs, info = super().reset(seed=seed, options=options)
+        obs = {k: v[:, 0] for k, v in obs.items()}
+        info = {k: v[:, 0] for k, v in info.items()}
+        return obs, info
+
+    def step(self, action: NDArray[np.floating]) -> tuple[dict, float, bool, bool, dict]:
+        obs, reward, terminated, truncated, info = super().step(action)
+        obs = {k: v[:, 0] for k, v in obs.items()}
+        info = {k: v[:, 0] for k, v in info.items()}
+        return obs, reward[:, 0], terminated[:, 0], truncated[:, 0], info