Clean up formatting

Author: amacati

benchmark/main.py

@@ -1,35 +1,36 @@
 import time
-from ml_collections import config_dict
 
+import gymnasium
 import jax
 import jax.numpy as jnp
 import numpy as np
-import gymnasium
+from ml_collections import config_dict
+
+import crazyflow  # noqa: F401, ensure gymnasium envs are registered
 from crazyflow.sim.core import Sim
-import crazyflow.gymnasium_envs
 
 
 def analyze_timings(times: list[float], n_steps: int, n_worlds: int, freq: float) -> None:
     """Analyze timing results and print performance metrics."""
     if not times:
         raise ValueError("The list of timing results is empty.")
-    
+
     tmin, idx_tmin = np.min(times), np.argmin(times)
     tmax, idx_tmax = np.max(times), np.argmax(times)
-    
+
     # Check for significant variance
     if tmax / tmin > 5:
         print("Warning: step time varies by more than 5x. Is JIT compiling during the benchmark?")
         print(f"Times: max {tmax:.2e}@{idx_tmax}, min {tmin:.2e}@{idx_tmin}")
-    
+
     # Performance metrics
     n_frames = n_steps * n_worlds  # Number of frames simulated
     total_time = np.sum(times)
     avg_step_time = np.mean(times)
     step_time_std = np.std(times)
     fps = n_frames / total_time
     real_time_factor = (n_steps / freq) * n_worlds / total_time
-    
+
     print(
         f"Avg step time: {avg_step_time:.2e}s, std: {step_time_std:.2e}"
         f"\nFPS: {fps:.3e}, Real time factor: {real_time_factor:.2e}"
@@ -62,7 +63,7 @@ def profile_gym_env_step(sim_config: config_dict.ConfigDict, n_steps: int, devic
     _, _ = envs.reset_all(seed=42)
     _, _, _, _, _ = envs.step(action)
     _, _ = envs.reset_all(seed=42)
-    
+
     jax.block_until_ready(envs.unwrapped.sim._mjx_data)  # Ensure JIT compiled dynamics
 
     # Step through the environment
@@ -85,7 +86,7 @@ def profile_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
 
     cmd = jnp.zeros((sim.n_worlds, sim.n_drones, 4), device=device)
     cmd = cmd.at[0, 0, 0].set(1)
-    
+
     sim.reset()
     sim.attitude_control(cmd)
     sim.step()

benchmark/performance.py

@@ -1,19 +1,16 @@
+import gymnasium
 import jax
-import jax.numpy as jnp
 import numpy as np
-import gymnasium
 from ml_collections import config_dict
 from pyinstrument import Profiler
 from pyinstrument.renderers.html import HTMLRenderer
 
+import crazyflow  # noqa: F401, ensure gymnasium envs are registered
 from crazyflow.sim.core import Sim
-import crazyflow.gymnasium_envs
 
 
 def profile_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
-    sim = Sim(
-        **sim_config
-    )
+    sim = Sim(**sim_config)
     device = jax.devices(device)[0]
     ndim = 13 if sim.control == "state" else 4
     control_fn = sim.state_control if sim.control == "state" else sim.attitude_control
@@ -39,6 +36,7 @@ def profile_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
     renderer = HTMLRenderer()
     renderer.open_in_browser(profiler.last_session)
 
+
 def profile_gym_env_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
     device = jax.devices(device)[0]
 
@@ -65,7 +63,7 @@ def profile_gym_env_step(sim_config: config_dict.ConfigDict, n_steps: int, devic
     _, _ = envs.reset_all(seed=42)
     _, _, _, _, _ = envs.step(action)
     _, _ = envs.reset_all(seed=42)
-    
+
     jax.block_until_ready(envs.unwrapped.sim._mjx_data)  # Ensure JIT compiled dynamics
 
     profiler = Profiler()
@@ -79,6 +77,7 @@ def profile_gym_env_step(sim_config: config_dict.ConfigDict, n_steps: int, devic
     renderer.open_in_browser(profiler.last_session)
     envs.close()
 
+
 def main():
     device = "cpu"
     sim_config = config_dict.ConfigDict()
@@ -104,6 +103,5 @@ def main():
     profile_gym_env_step(sim_config, 1000, device)
 
 
-
 if __name__ == "__main__":
     main()

crazyflow/gymnasium_envs/__init__.py

@@ -1,6 +1,3 @@
-from crazyflow.gymnasium_envs.crazyflow import CrazyflowEnvReachGoal, CrazyflowEnvTargetVelocity
-
-
 from gymnasium.envs.registration import register
 
 register(

crazyflow/gymnasium_envs/crazyflow.py

@@ -1,45 +1,38 @@
 import math
 import warnings
-from dataclasses import fields
 from functools import partial
 from typing import Dict, Literal, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
 import numpy as np
+from flax.struct import dataclass
 from gymnasium import spaces
 from gymnasium.vector import VectorEnv
 from gymnasium.vector.utils import batch_space
 from jax import Array
 
-from crazyflow.control.controller import Control, MAX_THRUST, MIN_THRUST
+from crazyflow.control.controller import MAX_THRUST, MIN_THRUST, Control
 from crazyflow.sim.core import Sim
 from crazyflow.sim.structs import SimState
-from flax.struct import dataclass
+
 
 @dataclass
 class RescaleParams:
     scale_factor: jnp.ndarray
     mean: jnp.ndarray
 
+
 CONTROL_RESCALE_PARAMS = {
-        "state": None, 
-        "thrust": None,
-        "attitude": RescaleParams(
-            scale_factor=jnp.array([
-                4 * (MAX_THRUST - MIN_THRUST) / 2, 
-                jnp.pi / 6, 
-                jnp.pi / 6, 
-                jnp.pi / 6
-            ]),
-            mean=jnp.array([
-                4 * (MIN_THRUST + MAX_THRUST) / 2, 
-                0.0, 
-                0.0, 
-                0.0
-            ])
+    "state": None,
+    "thrust": None,
+    "attitude": RescaleParams(
+        scale_factor=jnp.array(
+            [4 * (MAX_THRUST - MIN_THRUST) / 2, jnp.pi / 6, jnp.pi / 6, jnp.pi / 6]
         ),
-    }
+        mean=jnp.array([4 * (MIN_THRUST + MAX_THRUST) / 2, 0.0, 0.0, 0.0]),
+    ),
+}
 
 
 class CrazyflowBaseEnv(VectorEnv):
@@ -48,7 +41,7 @@ class CrazyflowBaseEnv(VectorEnv):
     def __init__(
         self,
         *,
-        jax_random_key,  # required for jax random number generator
+        jax_random_key: int,  # required for jax random number generator
         num_envs: int = 1,  # required for VectorEnv
         max_episode_steps: int = 1000,
         return_datatype: Literal["numpy", "jax"] = "jax",
@@ -57,9 +50,13 @@ def __init__(
         """Summary: Initializes the CrazyflowEnv.
 
         Args:
-        max_episode_steps (int): The maximum number of steps per episode.
-            return_datatype (Literal["numpy", "jax"]): The data type for returned arrays, either "numpy" or "jax". If specified as "numpy", the returned arrays will be numpy arrays on the CPU. If specified as "jax", the returned arrays will be jax arrays on the "device" specifiedf or the simulation.
-            **kwargs: Takes arguments that are passed to the Crazyfly simulation .
+            jax_random_key: The random key for the jax random number generator.
+            num_envs: The number of environments to run in parallel.
+            max_episode_steps: The maximum number of steps per episode.
+            return_datatype: The data type for returned arrays, either "numpy" or "jax". If "numpy",
+                the returned arrays will be numpy arrays on the CPU. If "jax", the returned arrays
+                will be jax arrays on the "device" specified for the simulation.
+            **kwargs: Takes arguments that are passed to the Crazyfly simulation.
         """
         assert num_envs == kwargs["n_worlds"], "num_envs must be equal to n_worlds"
 
@@ -167,7 +164,9 @@ def _rescale_action(action: Array, control_type: str) -> Array:
 
         return action * params.scale_factor + params.mean
 
-    def reset_all(self, *, seed: Optional[int] = None, options: Optional[dict] = None):
+    def reset_all(
+        self, *, seed: Optional[int] = None, options: Optional[dict] = None
+    ) -> tuple[dict[str, Array], dict]:
         super().reset(seed=seed)
 
         # Resets ALL (!) environments
@@ -209,15 +208,15 @@ def reset(self, mask: Array) -> None:
         )
 
     @property
-    def reward(self):
+    def reward(self) -> Array:
         return self._reward(self.terminated, self.sim.states)
 
     @property
-    def terminated(self):
+    def terminated(self) -> Array:
         return self._terminated(self.prev_done, self.sim.states, self.sim.contacts())
 
     @property
-    def truncated(self):
+    def truncated(self) -> Array:
         return self._truncated(
             self.prev_done, self.sim.steps, self.max_episode_steps, self.n_substeps
         )
@@ -250,7 +249,9 @@ def render(self):
     def _get_obs(self) -> Dict[str, jnp.ndarray]:
         obs = {
             state: self._maybe_to_numpy(
-                getattr(self.sim.states, state)[..., 2] if state == "pos" else getattr(self.sim.states, state)
+                getattr(self.sim.states, state)[..., 2]
+                if state == "pos"
+                else getattr(self.sim.states, state)
             )
             for state in self.states_to_include_in_obs
         }
@@ -278,7 +279,7 @@ def __init__(self, **kwargs: dict):
         self.goal = jnp.zeros((kwargs["n_worlds"], 3), dtype=jnp.float32)
 
     @property
-    def reward(self):
+    def reward(self) -> Array:
         return self._reward(self.terminated, self.sim.states, self.goal)
 
     @staticmethod
@@ -323,7 +324,7 @@ def __init__(self, **kwargs: dict):
         self.target_vel = jnp.zeros((kwargs["n_worlds"], 3), dtype=jnp.float32)
 
     @property
-    def reward(self):
+    def reward(self) -> Array:
         return self._reward(self.terminated, self.sim.states, self.target_vel)
 
     @staticmethod

examples/gradient.py

@@ -11,7 +11,7 @@
 def main():
     sim = Sim(control=Control.state)
 
-    def step(cmd: NDArray):
+    def step(cmd: NDArray) -> jax.Array:
         sim.reset()
         sim.state_control(cmd)
         sim.step()

examples/gymnasium_env.py

@@ -2,7 +2,6 @@
 import numpy as np
 from ml_collections import config_dict
 
-import crazyflow.gymnasium_envs
 from crazyflow.control.controller import Control, Controller
 from crazyflow.sim.physics import Physics
 
@@ -13,23 +12,24 @@
 sim_config.control = Control.default
 sim_config.controller = Controller.default
 sim_config.control_freq = 50
-sim_config.n_drones=1
-sim_config.n_worlds=20
+sim_config.n_drones = 1
+sim_config.n_worlds = 20
 
-SEED=42
+SEED = 42
 
 envs = gymnasium.make_vec(
     "CrazyflowEnvReachGoal-v0",
     max_episode_steps=1000,
     return_datatype="numpy",
-    num_envs=sim_config.n_worlds, 
+    num_envs=sim_config.n_worlds,
     jax_random_key=SEED,
     **sim_config,
 )
 
 # action for going up (in attitude control). NOTE actions are rescaled in the environment
 action = np.array(
-    [[[-0.2, 0, 0, 0] for _ in range(sim_config.n_drones)] for _ in range(sim_config.n_worlds)], dtype=np.float32
+    [[[-0.2, 0, 0, 0] for _ in range(sim_config.n_drones)] for _ in range(sim_config.n_worlds)],
+    dtype=np.float32,
 ).reshape(sim_config.n_worlds, -1)
 
 obs, info = envs.reset_all(seed=SEED)
@@ -38,5 +38,5 @@
 for _ in range(1500):
     observation, reward, terminated, truncated, info = envs.step(action)
     envs.render()
-    
+
 envs.close()

examples/spiral.py

@@ -1,21 +1,11 @@
 import numpy as np
 
-from crazyflow.control.controller import Control, Controller
+from crazyflow.control.controller import Control
 from crazyflow.sim.core import Sim
-from crazyflow.sim.physics import Physics
 
 
 def main():
-    sim = Sim(
-        n_worlds=1,
-        n_drones=4,
-        physics=Physics.analytical,
-        control=Control.state,
-        controller=Controller.emulatefirmware,
-        freq=500,
-        control_freq=500,
-        device="cpu",
-    )
+    sim = Sim(n_worlds=1, n_drones=4, control=Control.state, device="cpu")
     sim.reset()
     duration = 5.0
     fps = 60

pyproject.toml

@@ -29,7 +29,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-test = ["pytest>=8.0.0"]
+test = ["pytest>=8.0.0", "pytest-cov"]
 gpu = ["jax[cuda12]"]
 
 [tool.setuptools.packages]
@@ -69,13 +69,16 @@ target-version = "py38"
 
 [tool.ruff.lint]
 select = ["E4", "E7", "E9", "F", "I", "D", "TCH", "ANN"]
-ignore = ["ANN101", "ANN401"]
+ignore = ["ANN401"]
 fixable = ["ALL"]
 unfixable = []
 
 [tool.ruff.lint.per-file-ignores]
-"benchmarks/*" = ["D100", "D103"]
+"benchmark/*" = ["D100", "D103"]
 "tests/*" = ["D100", "D103", "D104"]
+"examples/*" = ["D100", "D103"]
+# TODO: Remove once everything is stable and document
+"crazyflow/*" = ["D100", "D101", "D102", "D104", "D107"]  
 
 
 [tool.ruff.lint.pydocstyle]