Add multi drone env

Author: amacati

benchmarks/config/test.toml

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import fire
 import numpy as np
-from sim import time_sim_attitude_step, time_sim_reset, time_sim_step
+from sim import time_multi_drone_reset, time_multi_drone_step, time_sim_reset, time_sim_step
 
 
 def print_benchmark_results(name: str, timings: list[float]):
@@ -12,16 +13,21 @@ def print_benchmark_results(name: str, timings: list[float]):
     print(f"FPS: {1 / np.mean(timings):.2f}")
 
 
-if __name__ == "__main__":
-    n_tests = 10
-    sim_steps = 10
-    timings = time_sim_reset(n_tests=n_tests)
+def main(n_tests: int = 10, sim_steps: int = 10, multi_drone: bool = False):
+    reset_fn, step_fn = time_sim_reset, time_sim_step
+    if multi_drone:
+        reset_fn, step_fn = time_multi_drone_reset, time_multi_drone_step
+    timings = reset_fn(n_tests=n_tests)
     print_benchmark_results(name="Sim reset", timings=timings)
-    timings = time_sim_step(n_tests=n_tests, sim_steps=sim_steps)
+    timings = step_fn(n_tests=n_tests, sim_steps=sim_steps)
     print_benchmark_results(name="Sim steps", timings=timings / sim_steps)
-    timings = time_sim_step(n_tests=n_tests, sim_steps=sim_steps, physics_mode="sys_id")
+    timings = step_fn(n_tests=n_tests, sim_steps=sim_steps, physics_mode="sys_id")
     print_benchmark_results(name="Sim steps (sys_id backend)", timings=timings / sim_steps)
-    timings = time_sim_step(n_tests=n_tests, sim_steps=sim_steps, physics_mode="mujoco")
+    timings = step_fn(n_tests=n_tests, sim_steps=sim_steps, physics_mode="mujoco")
     print_benchmark_results(name="Sim steps (mujoco backend)", timings=timings / sim_steps)
-    timings = time_sim_attitude_step(n_tests=n_tests, sim_steps=sim_steps)
+    timings = step_fn(n_tests=n_tests, sim_steps=sim_steps, physics_mode="sys_id")
     print_benchmark_results(name="Sim steps (sys_id backend)", timings=timings / sim_steps)
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

benchmarks/main.py

@@ -14,25 +14,75 @@
 
 from lsy_drone_racing.utils import load_config
 
-config = load_config(Path('{Path(__file__).parents[1] / "config/level0.toml"}'))
-
+config = load_config(Path('{Path(__file__).parents[1] / "config/level3.toml"}'))
 """
+
 env_setup_code = """
 import gymnasium
 
 import lsy_drone_racing
 
-env = gymnasium.make('DroneRacing-v0', config=config)
+env = gymnasium.make(
+    config.env.id,
+    freq=config.env.freq,
+    sim_config=config.sim,
+    sensor_range=config.env.sensor_range,
+    track=config.env.track,
+    disturbances=config.env.get("disturbances"),
+    randomizations=config.env.get("randomizations"),
+    random_resets=config.env.random_resets,
+    seed=config.env.seed,
+)
 env.reset()
 env.step(env.action_space.sample())  # JIT compile
 env.reset()
 """
+
 attitude_env_setup_code = """
 import gymnasium
 
 import lsy_drone_racing
 
-env = gymnasium.make('DroneRacingAttitude-v0', config=config)
+env = gymnasium.make('DroneRacingAttitude-v0',
+    config.env.id,
+    freq=config.env.freq,
+    sim_config=config.sim,
+    sensor_range=config.env.sensor_range,
+    track=config.env.track,
+    disturbances=config.env.get("disturbances"),
+    randomizations=config.env.get("randomizations"),
+    random_resets=config.env.random_resets,
+    seed=config.env.seed,
+)
+env.reset()
+env.step(env.action_space.sample())  # JIT compile
+env.reset()
+"""
+
+load_multi_drone_config_code = f"""
+from pathlib import Path
+
+from lsy_drone_racing.utils import load_config
+
+config = load_config(Path('{Path(__file__).parents[1] / "config/multi_level0.toml"}'))
+"""
+
+multi_drone_env_setup_code = """
+import gymnasium
+
+import lsy_drone_racing
+
+env = gymnasium.make('MultiDroneRacing-v0',
+    n_drones=config.env.n_drones,
+    freq=config.env.freq,
+    sim_config=config.sim,
+    sensor_range=config.env.sensor_range,
+    track=config.env.track,
+    disturbances=config.env.get("disturbances"),
+    randomizations=config.env.get("randomizations"),
+    random_resets=config.env.random_resets,
+    seed=config.env.seed,
+)
 env.reset()
 env.step(env.action_space.sample())  # JIT compile
 env.reset()
@@ -62,3 +112,25 @@ def time_sim_attitude_step(n_tests: int = 10, sim_steps: int = 100) -> NDArray[n
 for _ in range({sim_steps}):
     env.step(env.action_space.sample())"""
     return np.array(timeit.repeat(stmt=stmt, setup=setup, number=1, repeat=n_tests))
+
+
+def time_multi_drone_reset(n_tests: int = 10) -> NDArray[np.floating]:
+    setup = load_multi_drone_config_code + multi_drone_env_setup_code + "\nenv.reset()"
+    stmt = """env.reset()"""
+    return np.array(timeit.repeat(stmt=stmt, setup=setup, number=1, repeat=n_tests))
+
+
+def time_multi_drone_step(
+    n_tests: int = 10, sim_steps: int = 100, physics_mode: str = "analytical"
+) -> NDArray[np.floating]:
+    modify_config_code = f"""config.sim.physics = '{physics_mode}'\n"""
+    setup = (
+        load_multi_drone_config_code
+        + modify_config_code
+        + multi_drone_env_setup_code
+        + "\nenv.reset()"
+    )
+    stmt = f"""
+for _ in range({sim_steps}):
+    env.step(env.action_space.sample())"""
+    return np.array(timeit.repeat(stmt=stmt, setup=setup, number=1, repeat=n_tests))

benchmarks/sim.py

@@ -0,0 +1,78 @@
+# Level 0
+
+# | Evaluation Scenario | Rand. Inertial Properties | Randomized Obstacles, Gates | Rand. Between Episodes | Notes             |
+# | :-----------------: | :-----------------------: | :-------------------------: | :--------------------: | :---------------: |
+# | `level0.toml`       | *No*                      | *No*                        | *No*                   | Perfect knowledge |
+[controller]
+file = "trajectory_controller.py" # Put your controller file name here. Specifying a controller as argument to scripts will override this setting. Controllers are located in `lsy_drone_racing/control/`
+
+[deploy]
+### Settings only relevant for deployment
+# Whether to check if gate and obstacle positions observed by vicon are within the limits defined down below.
+check_race_track = true
+# Whether to check if the drone start position is within the limits specified down below.
+check_drone_start_pos = true
+# Lets you practice your controller without putting up gates & obstacles, assumes nominal positions given below.
+practice_without_track_objects = false
+
+[sim]
+# Physics options:
+# "analytical": Analytical, simplified dynamics model
+# "mujoco": Mujoco dynamics. May take longer to compile at startup.
+# "sys_id": System identification model.
+physics = "analytical"
+
+camera_view = [5.0, -40.0, -40.0, 0.5, -1.0, 0.5]
+freq = 500                                        # Simulation frequency, in Hz
+attitude_freq = 500                               # Controller frequency, in Hz. This frequency is used to simulate the onboard controller, NOT for the environment's step function
+gui = false                                       # Enable/disable PyBullet's GUI
+
+[sim.disturbances.action]
+type = "GaussianNoise"
+std = 0.001
+
+[sim.disturbances.dynamics]
+type = "UniformNoise"
+low = [-0.1, -0.1, -0.1]
+high = [0.1, 0.1, 0.1]
+
+[env]
+id = "MultiDroneRacing-v0" # Either "MultiDroneRacingEnv-v0" or "MultiDroneRacingAttitude-v0". If using "MultiDroneRacingAttitude-v0", the drone will use the attitude controller instead of the position controller.
+n_drones = 2          # Number of drones has to match the track configuration for drones
+random_resets = false # Whether to re-seed the random number generator between episodes
+seed = 1337           # Random seed
+freq = 50             # Frequency of the environment's step function, in Hz
+symbolic = false      # Whether to include symbolic expressions in the info dict. Note: This can interfere with multiprocessing! If you want to parallelize your training, set this to false.
+sensor_range = 0.45   # Range at which the exact location of gates and obstacles become visible to the drone. Objects that are not in the drone's sensor range report their nominal position.
+
+[env.track]
+# Tall gates: 1.0m height. Short gates: 0.525m height. Height is measured from the ground to the
+# center of the gate.
+[[env.track.gates]]
+pos = [0.45, -1.0, 0.56]
+rpy = [0.0, 0.0, 2.35]
+[[env.track.gates]]
+pos = [1.0, -1.55, 1.11]
+rpy = [0.0, 0.0, -0.78]
+[[env.track.gates]]
+pos = [0.0, 0.5, 0.56]
+rpy = [0.0, 0.0, 0.0]
+[[env.track.gates]]
+pos = [-0.5, -0.5, 1.11]
+rpy = [0.0, 0.0, 3.14]
+
+# Obstacle height: 1.4m. Height is measured from the ground to the top of the obstacle.
+[[env.track.obstacles]]
+pos = [1.0, -0.5, 1.4]
+[[env.track.obstacles]]
+pos = [0.5, -1.5, 1.4]
+[[env.track.obstacles]]
+pos = [-0.5, 0.0, 1.4]
+[[env.track.obstacles]]
+pos = [0.0, 1.0, 1.4]
+
+[env.track.drone]
+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]]

config/multi_level0.toml

@@ -46,3 +46,12 @@
     max_episode_steps=1800,
     disable_env_checker=True,
 )
+
+# region MultiEnvs
+
+register(
+    id="MultiDroneRacing-v0",
+    entry_point="lsy_drone_racing.envs.multi_drone_race:MultiDroneRacingEnv",
+    max_episode_steps=1800,
+    disable_env_checker=True,
+)