crazyflow 0.0.1__tar.gz

crazyflow-0.0.1/LICENSE

@@ -0,0 +1,9 @@
+The MIT License (MIT)
+Copyright © 2025 Martin Schuck
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+

crazyflow-0.0.1/PKG-INFO

@@ -0,0 +1,156 @@
+Metadata-Version: 2.4
+Name: crazyflow
+Version: 0.0.1
+Summary: Fast, parallelizable simulations of Crazyflies with JAX and MuJoCo.
+License: The MIT License (MIT)
+        Copyright © 2025 Martin Schuck
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+        
+        
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax!=0.5.3,>=0.5.0
+Requires-Dist: mujoco>=3.3.0
+Requires-Dist: mujoco-mjx>=3.3.0
+Requires-Dist: gymnasium
+Requires-Dist: imageio
+Requires-Dist: einops
+Requires-Dist: flax
+Requires-Dist: ml_collections
+Requires-Dist: casadi
+Requires-Dist: numpy
+Provides-Extra: test
+Requires-Dist: pytest>=8.0.0; extra == "test"
+Requires-Dist: pytest-cov; extra == "test"
+Requires-Dist: pytest-timeout; extra == "test"
+Provides-Extra: gpu
+Requires-Dist: jax[cuda12]; extra == "gpu"
+Provides-Extra: benchmark
+Requires-Dist: fire; extra == "benchmark"
+Dynamic: license-file
+
+![Crazyflow Logo](https://github.com/utiasDSL/crazyflow/raw/main/docs/img/logo.png)
+
+--------------------------------------------------------------------------------
+
+Fast, parallelizable simulations of Crazyflies with JAX and MuJoCo.
+
+[![Python Version]][Python Version URL] [![Ruff Check]][Ruff Check URL] [![Documentation Status]][Documentation Status URL] [![Tests]][Tests URL]
+
+[Python Version]: https://img.shields.io/badge/python-3.10+-blue.svg
+[Python Version URL]: https://www.python.org
+
+[Ruff Check]: https://github.com/utiasDSL/crazyflow/actions/workflows/ruff.yml/badge.svg?style=flat-square
+[Ruff Check URL]: https://github.com/utiasDSL/crazyflow/actions/workflows/ruff.yml
+
+[Documentation Status]: https://readthedocs.org/projects/crazyflow/badge/?version=latest
+[Documentation Status URL]: https://crazyflow.readthedocs.io/en/latest/?badge=latest
+
+[Tests]: https://github.com/utiasDSL/crazyflow/actions/workflows/testing.yml/badge.svg
+[Tests URL]: https://github.com/utiasDSL/crazyflow/actions/workflows/testing.yml
+
+
+## Architecture
+
+Crazyflow is a high-performance simulation framework for Crazyflie drones that leverages JAX for efficient parallelization and automatic differentiation. The architecture is designed around a flexible pipeline that can be configured at initialization time, enabling users to swap out physics backends, control methods, and integration schemes.
+
+### Core Components
+
+#### Simulation Pipeline
+The simulation is built as a pipeline of functions that are composed at initialization time based on the configuration. This approach avoids runtime branching and allows JAX to optimize the entire pipeline as a single computation. Users can insert their own pure functions into the pipeline to modify the simulation behavior while maintaining compatibility with JAX's optimizations.
+
+#### Physics Backends
+Multiple physics models are supported:
+- analytical: A first-principles model based on physical equations
+- sys_id: A system-identified model trained on real drone data
+- mujoco: MuJoCo physics engine for more complex interactions
+
+#### Control Modes
+Different control interfaces are available:
+- state: High-level control of position, velocity, and yaw
+- attitude: Mid-level control of collective thrust and orientation
+- thrust: Low-level control of individual motor thrusts
+
+#### Integration Methods
+For analytical and system-identified physics:
+- euler: Simple first-order integration
+- rk4: Fourth-order Runge-Kutta integration for higher accuracy
+
+### Parallelization
+Crazyflow supports massive parallelization across:
+- Worlds: Independent simulation environments that can run in parallel
+- Drones: Multiple drones within each world
+- Devices: Computations can be executed on CPU or GPU
+This parallelization is achieved through JAX's vectorization capabilities, allowing thousands of simulations to run simultaneously with minimal overhead.
+
+### Domain Randomization
+The framework supports domain randomization through the crazyflow/randomize module, allowing parameters like mass to be varied across simulations to improve sim-to-real transfer.
+
+### Functional Design
+The simulation follows a functional programming paradigm: All state is contained in immutable data structures. Updates create new states rather than modifying existing ones. All functions are pure, enabling JAX's transformations (JIT, grad, vmap) and thus automatic differentiation through the entire simulation, making it suitable for gradient-based optimization and reinforcement learning.
+
+## Examples
+The repository includes several example scripts demonstrating different capabilities:
+| Example                                   | Description                                                 |
+| ----------------------------------------- | ----------------------------------------------------------- |
+| [`hover.py`](examples/hover.py)           | Basic hovering using state control                          |
+| [`thrust.py`](examples/thrust.py)         | Direct motor control using thrust commands                  |
+| [`render.py`](examples/render.py)         | Visualization of multiple drones with motion traces         |
+| [`contacts.py`](examples/contacts.py)     | Collision detection between drones                          |
+| [`gradient.py`](examples/gradient.py)     | Computing gradients through the simulation for optimization |
+| [`change_pos.py`](examples/change_pos.py) | Manipulating drone positions programmatically               |
+
+
+
+## Known Issues
+- `"RuntimeError: MUJOCO_PATH environment variable is not set"` upon installing this package. This error can be resolved by using `venv` instead of `conda`. Somtimes the `mujoco` install can [fail with `conda`](https://github.com/google-deepmind/mujoco/issues/1004).
+- If using `zsh` don't forget to escape brackets when installing additional dependencies: `pip install .\[gpu\]`.
+
+### Using the project with VSCode devcontainers
+
+**Running on CPU**: by default the containers run on CPU. You don't need to take any action.
+
+**Running on GPU**: The devcontsainers can easily run using your computer's NVIDIA GPU on Linux and Windows. This makes sense if you want to accelerate simulation by running thousands of simulation in parallel. In order to work you need to install the [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=WSL-Ubuntu&target_version=2.0&target_type=deb_local), [NVIDIA Container runtime](https://developer.nvidia.com/container-runtime) for your computer. Finally, enable GPU access to the devcontainers by setting the commented out `"--gpus=all"` and `"--runtime=nvidia"` flags in `devcontainer.json`. 
+
+
+**Linux**
+1. Make sure to be in a X11 session ([link](https://askubuntu.com/questions/1410256/how-do-i-use-the-x-window-manager-instead-of-wayland-on-ubuntu-22-04)), otherwise rendering of the drone will fail.
+2. Install [Docker](https://docs.docker.com/engine/install/) (, and make sure Docker Daemon is running)
+3. Install [VSCode](https://code.visualstudio.com/), with [devcontainer extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers), and [remote dev pack](https://marketplace.visualstudio.com/items?itemName=ms-azuretools.vscode-docker).
+4. Clone this project's code. Rename `/.devcontainer/devcontainer.linux.json` to `/.devcontainer/devcontainer.json`.
+5. Open this project in VSCode. VSCode should automatically detect the devcontainer and prompt you to `Reopen in container`. If not, see [here](https://code.visualstudio.com/docs/devcontainers/containers#_quick-start-open-an-existing-folder-in-a-container) to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.
+
+**Windows** (requires Windows 10 or later)
+
+For windows, we require WSL2 to run the devcontainer. (So its actually Linux with extra steps.) Full instructions can be found [in the official docs](https://code.visualstudio.com/blogs/2020/07/01/containers-wsl#_getting-started). Here are the important steps:
+1. Install [Docker](https://docs.docker.com/desktop/setup/install/windows-install/), and WSL2, and Ubuntu 22.04 LTS (, and make sure Docker Daemon is running)
+2. Docker will recognize that you have WSL installed and prompt you via Windows Notifications to enable WSL integration -> confirm this with `Enable WSL integration`. If not, open `Docker Desktop`, navigate to the settings, and manually enable WSL integration. (There are TWO setting options for this. Make sure to enable BOTH!)
+3. Install [VSCode](https://code.visualstudio.com/), with the [WSL extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-wsl), [devcontainer extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers), and [remote dev pack](https://marketplace.visualstudio.com/items?itemName=ms-azuretools.vscode-docker).
+4. Clone the source code for the exercises in the WSL2 file system to `/home` (`~`), or wherever you like. (Performance when working on the WSL file system is much better compared to Windows file system). You can access the WSL filesystem either by starting a WSL/Ubuntu console, or via the Windows File Explorer at `\\wsl.localhost\Ubuntu\home` (replace `Ubuntu` with your distro, if necessary).
+7. Rename `/.devcontainer/devcontainer.windows.json` to `/.devcontainer/devcontainer.json`.
+8. Open this project in VSCode. The easiest way to do so is by starting a WSL/Ubuntu shell, navigating via `cd` to the source code, then type `code .` to open VSCode. VSCode should automatically detect the devcontainer and prompt you to `Reopen in container`. If not, see [here](https://code.visualstudio.com/docs/devcontainers/containers#_quick-start-open-an-existing-folder-in-a-container) to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.
+
+
+**MacOS**
+
+Unfortunately, we did not get the devcontainer to work with MacOS yet, even after following [those](https://gist.github.com/sorny/969fe55d85c9b0035b0109a31cbcb088) steps. We expect that the issue is related to Mujoco rendering from inside the Docker container and display forwarding with X11. There is also an [unresolved Issue](https://github.com/google-deepmind/mujoco/issues/1047) on GitHub. If you manage to make it work, please let us know.
+
+Until then, MacOS users are required to install this project using an python environment manager such as [conda](https://docs.anaconda.com/anaconda/install/) or [mamba](https://mamba.readthedocs.io/en/latest/). If you use conda, these are the required commands: ```conda create --name crazyflow -c conda-forge python=3.11```, ```conda activate crazyflow```, ```conda install pip```, ```pip install -e .```.
+
+____________
+
+Known Issues:
+   - if building docker container fails at `RUN apt-get update`, make sure your host systems time is set correct: https://askubuntu.com/questions/1511514/docker-build-fails-at-run-apt-update-error-failed-to-solve-process-bin-sh
+

crazyflow-0.0.1/README.md

@@ -0,0 +1,113 @@
+![Crazyflow Logo](https://github.com/utiasDSL/crazyflow/raw/main/docs/img/logo.png)
+
+--------------------------------------------------------------------------------
+
+Fast, parallelizable simulations of Crazyflies with JAX and MuJoCo.
+
+[![Python Version]][Python Version URL] [![Ruff Check]][Ruff Check URL] [![Documentation Status]][Documentation Status URL] [![Tests]][Tests URL]
+
+[Python Version]: https://img.shields.io/badge/python-3.10+-blue.svg
+[Python Version URL]: https://www.python.org
+
+[Ruff Check]: https://github.com/utiasDSL/crazyflow/actions/workflows/ruff.yml/badge.svg?style=flat-square
+[Ruff Check URL]: https://github.com/utiasDSL/crazyflow/actions/workflows/ruff.yml
+
+[Documentation Status]: https://readthedocs.org/projects/crazyflow/badge/?version=latest
+[Documentation Status URL]: https://crazyflow.readthedocs.io/en/latest/?badge=latest
+
+[Tests]: https://github.com/utiasDSL/crazyflow/actions/workflows/testing.yml/badge.svg
+[Tests URL]: https://github.com/utiasDSL/crazyflow/actions/workflows/testing.yml
+
+
+## Architecture
+
+Crazyflow is a high-performance simulation framework for Crazyflie drones that leverages JAX for efficient parallelization and automatic differentiation. The architecture is designed around a flexible pipeline that can be configured at initialization time, enabling users to swap out physics backends, control methods, and integration schemes.
+
+### Core Components
+
+#### Simulation Pipeline
+The simulation is built as a pipeline of functions that are composed at initialization time based on the configuration. This approach avoids runtime branching and allows JAX to optimize the entire pipeline as a single computation. Users can insert their own pure functions into the pipeline to modify the simulation behavior while maintaining compatibility with JAX's optimizations.
+
+#### Physics Backends
+Multiple physics models are supported:
+- analytical: A first-principles model based on physical equations
+- sys_id: A system-identified model trained on real drone data
+- mujoco: MuJoCo physics engine for more complex interactions
+
+#### Control Modes
+Different control interfaces are available:
+- state: High-level control of position, velocity, and yaw
+- attitude: Mid-level control of collective thrust and orientation
+- thrust: Low-level control of individual motor thrusts
+
+#### Integration Methods
+For analytical and system-identified physics:
+- euler: Simple first-order integration
+- rk4: Fourth-order Runge-Kutta integration for higher accuracy
+
+### Parallelization
+Crazyflow supports massive parallelization across:
+- Worlds: Independent simulation environments that can run in parallel
+- Drones: Multiple drones within each world
+- Devices: Computations can be executed on CPU or GPU
+This parallelization is achieved through JAX's vectorization capabilities, allowing thousands of simulations to run simultaneously with minimal overhead.
+
+### Domain Randomization
+The framework supports domain randomization through the crazyflow/randomize module, allowing parameters like mass to be varied across simulations to improve sim-to-real transfer.
+
+### Functional Design
+The simulation follows a functional programming paradigm: All state is contained in immutable data structures. Updates create new states rather than modifying existing ones. All functions are pure, enabling JAX's transformations (JIT, grad, vmap) and thus automatic differentiation through the entire simulation, making it suitable for gradient-based optimization and reinforcement learning.
+
+## Examples
+The repository includes several example scripts demonstrating different capabilities:
+| Example                                   | Description                                                 |
+| ----------------------------------------- | ----------------------------------------------------------- |
+| [`hover.py`](examples/hover.py)           | Basic hovering using state control                          |
+| [`thrust.py`](examples/thrust.py)         | Direct motor control using thrust commands                  |
+| [`render.py`](examples/render.py)         | Visualization of multiple drones with motion traces         |
+| [`contacts.py`](examples/contacts.py)     | Collision detection between drones                          |
+| [`gradient.py`](examples/gradient.py)     | Computing gradients through the simulation for optimization |
+| [`change_pos.py`](examples/change_pos.py) | Manipulating drone positions programmatically               |
+
+
+
+## Known Issues
+- `"RuntimeError: MUJOCO_PATH environment variable is not set"` upon installing this package. This error can be resolved by using `venv` instead of `conda`. Somtimes the `mujoco` install can [fail with `conda`](https://github.com/google-deepmind/mujoco/issues/1004).
+- If using `zsh` don't forget to escape brackets when installing additional dependencies: `pip install .\[gpu\]`.
+
+### Using the project with VSCode devcontainers
+
+**Running on CPU**: by default the containers run on CPU. You don't need to take any action.
+
+**Running on GPU**: The devcontsainers can easily run using your computer's NVIDIA GPU on Linux and Windows. This makes sense if you want to accelerate simulation by running thousands of simulation in parallel. In order to work you need to install the [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=WSL-Ubuntu&target_version=2.0&target_type=deb_local), [NVIDIA Container runtime](https://developer.nvidia.com/container-runtime) for your computer. Finally, enable GPU access to the devcontainers by setting the commented out `"--gpus=all"` and `"--runtime=nvidia"` flags in `devcontainer.json`. 
+
+
+**Linux**
+1. Make sure to be in a X11 session ([link](https://askubuntu.com/questions/1410256/how-do-i-use-the-x-window-manager-instead-of-wayland-on-ubuntu-22-04)), otherwise rendering of the drone will fail.
+2. Install [Docker](https://docs.docker.com/engine/install/) (, and make sure Docker Daemon is running)
+3. Install [VSCode](https://code.visualstudio.com/), with [devcontainer extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers), and [remote dev pack](https://marketplace.visualstudio.com/items?itemName=ms-azuretools.vscode-docker).
+4. Clone this project's code. Rename `/.devcontainer/devcontainer.linux.json` to `/.devcontainer/devcontainer.json`.
+5. Open this project in VSCode. VSCode should automatically detect the devcontainer and prompt you to `Reopen in container`. If not, see [here](https://code.visualstudio.com/docs/devcontainers/containers#_quick-start-open-an-existing-folder-in-a-container) to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.
+
+**Windows** (requires Windows 10 or later)
+
+For windows, we require WSL2 to run the devcontainer. (So its actually Linux with extra steps.) Full instructions can be found [in the official docs](https://code.visualstudio.com/blogs/2020/07/01/containers-wsl#_getting-started). Here are the important steps:
+1. Install [Docker](https://docs.docker.com/desktop/setup/install/windows-install/), and WSL2, and Ubuntu 22.04 LTS (, and make sure Docker Daemon is running)
+2. Docker will recognize that you have WSL installed and prompt you via Windows Notifications to enable WSL integration -> confirm this with `Enable WSL integration`. If not, open `Docker Desktop`, navigate to the settings, and manually enable WSL integration. (There are TWO setting options for this. Make sure to enable BOTH!)
+3. Install [VSCode](https://code.visualstudio.com/), with the [WSL extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-wsl), [devcontainer extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers), and [remote dev pack](https://marketplace.visualstudio.com/items?itemName=ms-azuretools.vscode-docker).
+4. Clone the source code for the exercises in the WSL2 file system to `/home` (`~`), or wherever you like. (Performance when working on the WSL file system is much better compared to Windows file system). You can access the WSL filesystem either by starting a WSL/Ubuntu console, or via the Windows File Explorer at `\\wsl.localhost\Ubuntu\home` (replace `Ubuntu` with your distro, if necessary).
+7. Rename `/.devcontainer/devcontainer.windows.json` to `/.devcontainer/devcontainer.json`.
+8. Open this project in VSCode. The easiest way to do so is by starting a WSL/Ubuntu shell, navigating via `cd` to the source code, then type `code .` to open VSCode. VSCode should automatically detect the devcontainer and prompt you to `Reopen in container`. If not, see [here](https://code.visualstudio.com/docs/devcontainers/containers#_quick-start-open-an-existing-folder-in-a-container) to open manually. Note: Opening the container for the first time might take a while (up to 15 min), as the container is pulled from the web and build.
+
+
+**MacOS**
+
+Unfortunately, we did not get the devcontainer to work with MacOS yet, even after following [those](https://gist.github.com/sorny/969fe55d85c9b0035b0109a31cbcb088) steps. We expect that the issue is related to Mujoco rendering from inside the Docker container and display forwarding with X11. There is also an [unresolved Issue](https://github.com/google-deepmind/mujoco/issues/1047) on GitHub. If you manage to make it work, please let us know.
+
+Until then, MacOS users are required to install this project using an python environment manager such as [conda](https://docs.anaconda.com/anaconda/install/) or [mamba](https://mamba.readthedocs.io/en/latest/). If you use conda, these are the required commands: ```conda create --name crazyflow -c conda-forge python=3.11```, ```conda activate crazyflow```, ```conda install pip```, ```pip install -e .```.
+
+____________
+
+Known Issues:
+   - if building docker container fails at `RUN apt-get update`, make sure your host systems time is set correct: https://askubuntu.com/questions/1511514/docker-build-fails-at-run-apt-update-error-failed-to-solve-process-bin-sh
+

crazyflow-0.0.1/benchmark/compile.py

@@ -0,0 +1,39 @@
+import time
+
+import fire
+import jax
+
+from crazyflow import Sim
+
+
+def main(cache: bool = False):
+    """Main entry point for profiling."""
+    if cache:
+        jax.config.update("jax_compilation_cache_dir", "/tmp/jax_cache")
+        jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
+        jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
+        jax.config.update("jax_persistent_cache_enable_xla_caches", "all")
+
+    # Time initialization
+    start = time.perf_counter()
+    sim = Sim(n_worlds=1, n_drones=1, physics="sys_id", control="attitude")
+    init_time = time.perf_counter() - start
+
+    # Time reset compilation
+    start = time.perf_counter()
+    sim._reset.lower(sim.data, sim.default_data, None).compile()
+    reset_time = time.perf_counter() - start
+
+    # Time step compilation
+    start = time.perf_counter()
+    sim._step.lower(sim.data, 1).compile()
+    step_time = time.perf_counter() - start
+
+    print(f"Simulation startup times | {sim.physics} | {sim.control}")
+    print(f"Initialization: {init_time:.2f}s")
+    print(f"Reset: {reset_time:.2f}s")
+    print(f"Step: {step_time:.2f}s")
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

crazyflow-0.0.1/benchmark/main.py

@@ -0,0 +1,156 @@
+import time
+
+import gymnasium
+import jax
+import jax.numpy as jnp
+import numpy as np
+from ml_collections import config_dict
+
+import crazyflow  # noqa: F401, ensure gymnasium envs are registered
+from crazyflow.sim import Sim
+
+
+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 > 10:
+        print("Warning: Fn time varies by more than 10x. 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 fn time: {avg_step_time:.2e}s, std: {step_time_std:.2e}"
+        f"\nFPS: {fps:.3e}, Real time factor: {real_time_factor:.2e}\n"
+    )
+
+
+def profile_gym_env_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
+    """Profile the Crazyflow gym environment step performance."""
+    times = []
+    device = jax.devices(device)[0]
+
+    envs = gymnasium.make_vec(
+        "DroneReachPos-v0",
+        time_horizon_in_seconds=3,
+        num_envs=sim_config.n_worlds,
+        device=sim_config.device,
+        freq=sim_config.attitude_freq,
+        physics=sim_config.physics,
+    )
+
+    # Action for going up (in attitude control)
+    action = np.zeros((sim_config.n_worlds, 4), dtype=np.float32)
+    action[..., 0] = 0.3
+    # Step through env once to ensure JIT compilation
+    envs.reset(seed=42)
+    envs.step(action)
+
+    jax.block_until_ready(envs.unwrapped.sim.data)  # Ensure JIT compiled dynamics
+
+    # Step through the environment
+    for _ in range(n_steps):
+        tstart = time.perf_counter()
+        envs.step(action)
+        jax.block_until_ready(envs.unwrapped.sim.data)
+        times.append(time.perf_counter() - tstart)
+
+    envs.close()
+    print("Gym env step performance:")
+    analyze_timings(times, n_steps, envs.unwrapped.sim.n_worlds, envs.unwrapped.sim.freq)
+
+
+def profile_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
+    """Profile the Crazyflow simulator step performance."""
+    sim = Sim(**sim_config)
+    times = []
+    device = jax.devices(device)[0]
+
+    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(sim.freq // sim.control_freq)
+    jax.block_until_ready(sim.data)  # Ensure JIT compiled dynamics
+
+    for _ in range(n_steps):
+        tstart = time.perf_counter()
+        sim.attitude_control(cmd)
+        sim.step(sim.freq // sim.control_freq)
+        jax.block_until_ready(sim.data)
+        times.append(time.perf_counter() - tstart)
+
+    print("Sim step performance:")
+    analyze_timings(times, n_steps, sim.n_worlds, sim.freq)
+
+
+def profile_reset(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
+    """Profile the Crazyflow simulator reset performance."""
+    sim = Sim(**sim_config)
+    times = []
+    times_masked = []
+    device = jax.devices(device)[0]
+
+    # Ensure JIT compiled reset
+    sim.reset()
+    jax.block_until_ready(sim.data)
+
+    # Test full reset
+    for _ in range(n_steps):
+        tstart = time.perf_counter()
+        sim.reset()
+        jax.block_until_ready(sim.data)
+        times.append(time.perf_counter() - tstart)
+
+    # Test masked reset (only reset first world)
+    mask = jnp.zeros(sim.n_worlds, dtype=bool, device=device)
+    mask = mask.at[0].set(True)
+    sim.reset(mask)
+    jax.block_until_ready(sim.data)
+
+    for _ in range(n_steps):
+        tstart = time.perf_counter()
+        sim.reset(mask)
+        jax.block_until_ready(sim.data)
+        times_masked.append(time.perf_counter() - tstart)
+
+    print("Sim reset performance:")
+    analyze_timings(times, n_steps, sim.n_worlds, sim.freq)
+    print("Sim masked reset performance:")
+    analyze_timings(times_masked, n_steps, sim.n_worlds, sim.freq)
+
+
+def main():
+    """Main entry point for profiling."""
+    device = "cpu"
+    sim_config = config_dict.ConfigDict()
+    sim_config.n_worlds = 1
+    sim_config.n_drones = 1
+    sim_config.physics = "analytical"
+    sim_config.control = "attitude"
+    sim_config.attitude_freq = 500
+    sim_config.device = device
+
+    print("Simulator performance\n")
+    profile_step(sim_config, 1000, device)
+    profile_reset(sim_config, 1000, device)
+
+    print("Gymnasium environment performance\n")
+    profile_gym_env_step(sim_config, 1000, device)
+
+
+if __name__ == "__main__":
+    main()

crazyflow-0.0.1/benchmark/op_count.py

@@ -0,0 +1,17 @@
+import crazyflow  # noqa: F401, ensure gymnasium envs are registered
+from crazyflow.sim import Sim
+
+
+def main():
+    """Main entry point for profiling."""
+    sim = Sim(n_worlds=1, n_drones=1, physics="analytical", control="attitude")
+
+    compiled_reset = sim._reset.lower(sim.data, sim.default_data, None).compile()
+    compiled_step = sim._step.lower(sim.data, 1).compile()
+    op_count_reset = compiled_reset.cost_analysis()["flops"]
+    op_count_step = compiled_step.cost_analysis()["flops"]
+    print(f"Op counts:\n Reset: {op_count_reset}\n Step: {op_count_step}")
+
+
+if __name__ == "__main__":
+    main()

crazyflow-0.0.1/benchmark/performance.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import gymnasium
+import jax
+import numpy as np
+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 import Sim
+
+if TYPE_CHECKING:
+    from crazyflow.gymnasium_envs import CrazyflowEnvReachGoal
+
+
+def profile_step(sim_config: config_dict.ConfigDict, n_steps: int, device: str):
+    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
+    cmd = np.zeros((sim.n_worlds, sim.n_drones, ndim))
+    # Ensure JIT compiled dynamics and control
+    sim.reset()
+    control_fn(cmd)
+    sim.step()
+    jax.block_until_ready(sim.data)
+
+    profiler = Profiler()
+    profiler.start()
+
+    for _ in range(n_steps):
+        control_fn(cmd)
+        # sim.reset()
+        sim.step()
+        jax.block_until_ready(sim.data)
+    profiler.stop()
+    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]
+
+    envs: CrazyflowEnvReachGoal = gymnasium.make_vec(
+        "DroneReachPos-v0", time_horizon_in_seconds=2, num_envs=sim_config.n_worlds, **sim_config
+    )
+
+    # Action for going up (in attitude control)
+    action = np.zeros((sim_config.n_worlds, 4), dtype=np.float32)
+    action[..., 0] = 0.3
+
+    # Step through env once to ensure JIT compilation.
+    envs.reset(seed=42)
+    envs.step(action)
+    envs.step(action)  # Ensure all paths have been taken at least once
+    envs.reset(seed=42)
+    jax.block_until_ready(envs.unwrapped.sim.data)
+
+    profiler = Profiler()
+    profiler.start()
+
+    for _ in range(n_steps):
+        envs.step(action)
+        jax.block_until_ready(envs.unwrapped.sim.data)
+
+    profiler.stop()
+    renderer = HTMLRenderer()
+    renderer.open_in_browser(profiler.last_session)
+    envs.close()
+
+
+def main():
+    device = "cpu"
+    sim_config = config_dict.ConfigDict()
+    sim_config.n_worlds = 1
+    sim_config.n_drones = 1
+    sim_config.physics = "analytical"
+    sim_config.control = "attitude"
+    sim_config.device = device
+
+    profile_step(sim_config, 1000, device)
+    profile_gym_env_step(sim_config, 1000, device)
+
+
+if __name__ == "__main__":
+    main()

crazyflow-0.0.1/crazyflow/__init__.py

@@ -0,0 +1,6 @@
+import crazyflow.gymnasium_envs  # noqa: F401, ensure gymnasium envs are registered
+from crazyflow.control import Control
+from crazyflow.sim import Physics, Sim
+
+__all__ = ["Sim", "Physics", "Control"]
+__version__ = "0.0.1"

crazyflow-0.0.1/crazyflow/constants.py

@@ -0,0 +1,13 @@
+import numpy as np
+from numpy.typing import NDArray
+
+# Physical constants
+GRAVITY: float = 9.81
+
+# Drone constants
+ARM_LEN: float = 0.0325 * np.sqrt(2)
+MIX_MATRIX: NDArray = np.array([[-0.5, -0.5, -1], [-0.5, 0.5, 1], [0.5, 0.5, -1], [0.5, -0.5, 1]])
+SIGN_MIX_MATRIX: NDArray = np.sign(MIX_MATRIX)
+MASS: float = 0.027
+J: NDArray = np.array([[2.3951e-5, 0, 0], [0, 2.3951e-5, 0], [0, 0, 3.2347e-5]])
+J_INV: NDArray = np.linalg.inv(J)

crazyflow-0.0.1/crazyflow/control/__init__.py

@@ -0,0 +1,3 @@
+from crazyflow.control.control import Control
+
+__all__ = ["Control"]